BelGio13/Qwen2.5-0.5B-sft-nl2sh

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# --- 1. Configuration ---
# Path to the base model on Hugging Face
BASE_MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"

# Path to the adapters on Hugging Face
ADAPTER_PATH = "BelGio13/Qwen2.5-0.5B-sft-nl2sh"



# --- 3. Load the Base Model and Tokenizer ---
print(f"Loading base model: {BASE_MODEL_ID}...")
# Load in bfloat16 for efficiency, just like in training
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID, torch_dtype=torch.bfloat16, trust_remote_code=True
)

print(f"Loading tokenizer for {BASE_MODEL_ID}...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token


# --- 4. Load and Apply the PEFT Adapters ---
print(f"Loading LoRA adapters from: {ADAPTER_PATH}...")
# The PeftModel class merges the adapters into the base model seamlessly.
model = PeftModel.from_pretrained(base_model, ADAPTER_PATH)
print("Adapters loaded and applied successfully.")


# --- 5. Inference Function ---
def generate_command(prompt: str, model, tokenizer):
    """
    Takes a user prompt and generates the nl2sh JSON output.
    """
    model.eval()  # Set the model to evaluation mode

    # Use the chat template to format the input correctly for the model
    messages = [{"role": "user", "content": prompt}]

    # Tokenize the input
    input_ids = tokenizer.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
    )

    # Generate the response
    with torch.no_grad():  # Disable gradient calculation for inference
        outputs = model.generate(
            input_ids=input_ids,
            max_new_tokens=256,  # Max length of the generated output
            eos_token_id=tokenizer.eos_token_id,
            do_sample=False,  # Use greedy decoding for deterministic output
            # For more creative but less consistent results, try:
            # do_sample=True,
            # temperature=0.6,
            # top_p=0.9,
        )

    # Decode the generated tokens, skipping the prompt tokens
    response_ids = outputs[0][input_ids.shape[-1] :]
    response_text = tokenizer.decode(response_ids, skip_special_tokens=True)

    return response_text


# --- 6. Interactive Command Loop ---
if __name__ == "__main__":
    print("\n--- NL2SH Fine-Tuned Model ---")
    print("Enter your command in natural language. Type 'exit' or 'quit' to end.")

    while True:
        try:
            user_prompt = input("\n> ")
            if user_prompt.lower() in ["exit", "quit"]:
                break

            if not user_prompt:
                continue

            # Generate the response from the fine-tuned model
            model_response = generate_command(user_prompt, model, tokenizer)

            print("\n🤖 Model Output:")
            print(model_response)

        except KeyboardInterrupt:
            print("\nExiting...")
            break

import torch
import wandb
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig

# -------------------------------------------------
# 1. Configuration
# -------------------------------------------------
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
DATASET_PATH = "./training_dataset.jsonl"
NEW_MODEL_NAME = "qwen2-0.5b-nl2sh-v1"
MAX_SEQ_LENGTH = 512

# -------------------------------------------------
# 2. Device (cuda)
# -------------------------------------------------
if not torch.cuda.is_available():
    raise RuntimeError("cuda not available.")
device = torch.device("cuda")
print(f"Using device: {device}")

# -------------------------------------------------
# 3. Load dataset
# -------------------------------------------------
dataset = load_dataset("json", data_files=DATASET_PATH, split="train").shuffle()
dataset = dataset.train_test_split(test_size=0.2)
train_set = dataset["train"]
eval_set = dataset["test"]

# -------------------------------------------------
# 4. Model & Tokenizer
# -------------------------------------------------
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)
model = model.to(device)
model.config.use_cache = False

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# -------------------------------------------------
# 5. LoRA Config
# -------------------------------------------------
peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=32,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
)

# -------------------------------------------------
# 7. SFTTrainer (trl 0.24.0 API)
# -------------------------------------------------


sft_config = SFTConfig(
    dataset_text_field="messages",
    max_length=MAX_SEQ_LENGTH,
    packing=False,
    output_dir=f"./{NEW_MODEL_NAME}-results",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    gradient_accumulation_steps=1,
    optim="adamw_torch",
    save_steps=50,
    logging_steps=10,
    learning_rate=2e-4,
    weight_decay=0.001,
    bf16=False,
    max_grad_norm=0.3,
    warmup_ratio=0.03,
    group_by_length=True,
    lr_scheduler_type="constant",
    report_to=["wandb"],
    project="nl2sh-sft",
    eval_strategy="steps",
    eval_steps=50,
    do_eval=True,
    logging_dir=f"./{NEW_MODEL_NAME}-logs",
)

run = wandb.init(
    project="nl2sh-sft",
    config=sft_config.to_dict(),  # Log the SFTConfig hyperparameters
)

trainer = SFTTrainer(
    model=model,
    args=sft_config,
    train_dataset=train_set,
    eval_dataset=eval_set,
    peft_config=peft_config,
    # These are the ONLY valid args in trl 0.24
)

print("Starting fine-tuning on MPS...")
trainer.train()
print("Fine-tuning complete.")
try:
    run.finish()
except Exception as e:
    print(e)