Skip to content

Models API

LegionHetoModel

Main model class for loading and fine-tuning.

from legionheto import LegionHetoModel

model = LegionHetoModel(
    model_name="meta-llama/Llama-2-7b-hf",
    max_seq_length=2048,
    load_in_4bit=True,
)

Methods

setup_lora

Configure LoRA for efficient fine-tuning.

model.setup_lora(
    r=16,
    alpha=32,
    dropout=0.05,
    target_modules=None,
)

save_adapter

Save LoRA adapter weights.

model.save_adapter("./adapter")

load_adapter

Load LoRA adapter weights.

model.load_adapter("./adapter")

merge_and_unload

Merge adapter with base model.

model.merge_and_unload()

generate

Generate text from prompt.

output = model.generate(
    prompt="Hello",
    max_new_tokens=128,
    temperature=0.7,
)

ModelRegistry

Registry for model architecture detection.

from legionheto.core import ModelRegistry

config = ModelRegistry.get_optimal_config("meta-llama/Llama-2-7b-hf")