scripts/mythos_lora_core.py

"""Shared LoRA train/test logic for local scripts and ZeroGPU Space."""

from __future__ import annotations

import json
from pathlib import Path

import torch
from datasets import Dataset
from peft import LoraConfig, PeftModel, TaskType, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer


BASE_MODEL = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
SYSTEM_PROMPT = (
    "You are Mythos-Coder, a coding agent that inspects the task, makes a "
    "concise plan, edits safely, verifies results, and explains fixes clearly."
)
DEFAULT_MAX_NEW_TOKENS = 256
LFS_POINTER_PREFIX = "version https://git-lfs.github.com/spec/v1"


def load_sft_jsonl(train_path: Path) -> Dataset:
    """Load JSONL SFT rows, with a clear error if Git LFS pointers were not pulled."""
    train_path = Path(train_path)
    if not train_path.exists():
        raise FileNotFoundError(f"Training file not found: {train_path}")

    rows: list[dict] = []
    with open(train_path, "r", encoding="utf-8") as handle:
        for line_num, line in enumerate(handle, 1):
            stripped = line.strip()
            if not stripped:
                continue
            if line_num == 1 and stripped.startswith(LFS_POINTER_PREFIX):
                raise ValueError(
                    f"{train_path} is a Git LFS pointer, not the actual dataset. "
                    "Run: git lfs install && git lfs pull"
                )
            try:
                row = json.loads(stripped)
            except json.JSONDecodeError as exc:
                raise ValueError(
                    f"Invalid JSON on line {line_num} of {train_path}: {exc}"
                ) from exc
            rows.append(row)

    if not rows:
        raise ValueError(f"No training examples found in {train_path}")

    return Dataset.from_list(rows)


def use_cuda() -> bool:
    return torch.cuda.is_available()


def load_tokenizer(model_name: str = BASE_MODEL):
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    return tokenizer


def train_lora(
    train_path: Path,
    output_dir: Path,
    *,
    model_name: str = BASE_MODEL,
    num_train_epochs: int = 1,
    per_device_train_batch_size: int = 2,
    gradient_accumulation_steps: int = 2,
    learning_rate: float = 1e-4,
    max_length: int = 2048,
) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    cuda = use_cuda()
    dtype = torch.float16 if cuda else torch.float32

    print(f"Base model:         {model_name}")
    print(f"Train file:         {train_path}")
    print(f"Output dir:         {output_dir}")
    print(f"Epochs:             {num_train_epochs}")
    print(f"Batch size:         {per_device_train_batch_size}")
    print(f"Grad accum:         {gradient_accumulation_steps}")
    print(f"Learning rate:      {learning_rate}")
    print(f"Max seq length:     {max_length}")
    print(f"fp16:               {cuda}")

    tokenizer = load_tokenizer(model_name)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=dtype,
        trust_remote_code=True,
    )
    model = model.to("cuda" if cuda else "cpu")

    lora_config = LoraConfig(
        r=8,
        lora_alpha=16,
        lora_dropout=0.05,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
        target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ],
    )
    model = get_peft_model(model, lora_config)

    dataset = load_sft_jsonl(train_path)
    print(f"Loaded examples:    {len(dataset)}")

    def formatting_func(example):
        if isinstance(example["messages"][0], list):
            return [
                tokenizer.apply_chat_template(messages, tokenize=False)
                for messages in example["messages"]
            ]
        return tokenizer.apply_chat_template(example["messages"], tokenize=False)

    training_args = SFTConfig(
        output_dir=str(output_dir),
        num_train_epochs=num_train_epochs,
        per_device_train_batch_size=per_device_train_batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        learning_rate=learning_rate,
        logging_steps=5,
        save_strategy="epoch",
        fp16=cuda,
        bf16=False,
        report_to="none",
        max_length=max_length,
        packing=False,
        dataset_text_field=None,
    )

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        processing_class=tokenizer,
        formatting_func=formatting_func,
    )
    train_result = trainer.train()
    trainer.save_model(str(output_dir))
    tokenizer.save_pretrained(str(output_dir))

    print(f"Training finished. LoRA adapter saved to: {output_dir}")

    return {
        "status": "ok",
        "model_name": model_name,
        "cuda": cuda,
        "train_examples": len(dataset),
        "output_dir": str(output_dir),
        "train_loss": float(train_result.training_loss) if train_result.training_loss else None,
        "runtime_seconds": float(train_result.metrics.get("train_runtime", 0)),
    }


def load_lora_model(model_path: Path, model_name: str = BASE_MODEL):
    cuda = use_cuda()
    dtype = torch.float16 if cuda else torch.float32
    device = "cuda" if cuda else "cpu"

    tokenizer = load_tokenizer(model_name)
    base_model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=dtype,
        trust_remote_code=True,
    )
    model = PeftModel.from_pretrained(base_model, str(model_path))
    model = model.to(device)
    model.eval()
    return model, tokenizer, device


def generate_response(
    model,
    tokenizer,
    prompt: str,
    device: str,
    *,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
) -> str:
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": prompt},
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = tokenizer(text, return_tensors="pt").to(device)
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    generated = output_ids[0][inputs["input_ids"].shape[1] :]
    return tokenizer.decode(generated, skip_special_tokens=True).strip()


def load_prompts(path: Path) -> list[str]:
    prompts = []
    with open(path, "r", encoding="utf-8") as handle:
        for line_num, line in enumerate(handle, 1):
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            prompt = str(row.get("prompt", "")).strip()
            if prompt:
                prompts.append(prompt)
    return prompts


def run_eval(
    prompts_path: Path,
    model_path: Path,
    output_path: Path,
    *,
    model_name: str = BASE_MODEL,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
) -> dict:
    prompts = load_prompts(prompts_path)
    if not prompts:
        raise ValueError(f"No prompts found in {prompts_path}")

    print(f"Base model:  {model_name}")
    print(f"LoRA path:   {model_path}")
    print(f"Prompts:     {len(prompts)} from {prompts_path}")

    model, tokenizer, device = load_lora_model(model_path, model_name=model_name)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    results = []
    for index, prompt in enumerate(prompts, 1):
        try:
            response = generate_response(
                model,
                tokenizer,
                prompt,
                device,
                max_new_tokens=max_new_tokens,
            )
        except Exception as exc:
            response = f"[generation error] {exc}"
        results.append(
            {
                "prompt": prompt,
                "generated_response": response,
                "max_new_tokens": max_new_tokens,
                "model_name": model_name,
                "model_path": str(model_path),
            }
        )

    with open(output_path, "w", encoding="utf-8") as handle:
        for row in results:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")

    return {
        "status": "ok",
        "cuda": use_cuda(),
        "prompt_count": len(results),
        "output_path": str(output_path),
    }