Upload code/finetune_qwen.py with huggingface_hub

code/finetune_qwen.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+"""Continue fine-tuning Qwen2.5-0.5B for category and company inference."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+
+PACKAGE_ROOT = Path(__file__).resolve().parent.parent
+if str(PACKAGE_ROOT) not in sys.path:
+    sys.path.insert(0, str(PACKAGE_ROOT))
+
+from pipeline.augment_training_data import sanitize_training_description
+from pipeline.company_inference import infer_company_name
+from pipeline.training_schema import CATEGORIES, INCOME_CATEGORIES, NON_INCOME_CATEGORIES
+
+MODEL_NAME = "Qwen/Qwen2.5-0.5B"
+DATA_PATH = PACKAGE_ROOT / "data" / "training_data.json"
+OUTPUT_DIR = PACKAGE_ROOT / "data" / "qwen-lora-adapter-0.5b"
+
+SYSTEM_PROMPT = (
+    "You are a bank transaction classifier for Indian bank statements. "
+    "Given a raw transaction description, infer both its category and the actual company when evidence exists. "
+    "Respond with ONLY a JSON object: "
+    '{"category": "<category>", "company_name": "<company_or_null>", "is_income": false, "confidence": 0.0}. '
+    f"Categories: {', '.join(CATEGORIES)}. "
+    "Use company_name=null for personal transfers or when the company is not supported by the description. "
+    "Credits to known employers = salary. UPI to person names = personal_transfer. "
+    "Refunds/reversals = original category. If truly unknown, category=unclassified, confidence=0.30."
+)
+
+
+def format_training_example(item: dict) -> dict[str, str]:
+    """Create one category + company prompt/completion training pair."""
+    description = item["description"]
+    category = item["category"]
+    if "is_income" in item:
+        is_income = bool(item["is_income"])
+    elif category in NON_INCOME_CATEGORIES:
+        is_income = False
+    else:
+        is_income = item.get("type") == "credit" or category in INCOME_CATEGORIES
+    company_name = infer_company_name(
+        description,
+        category=category,
+        explicit_name=item.get("company_name") or item.get("merchant") or item.get("counterparty"),
+    )
+    sanitized_description = sanitize_training_description(
+        description,
+        category=category,
+        company_name=company_name,
+    )
+    prompt = f"### System:\n{SYSTEM_PROMPT}\n\n### Input:\n{sanitized_description}\n\n### Output:\n"
+    completion = json.dumps({
+        "category": category,
+        "company_name": company_name,
+        "is_income": is_income,
+        "confidence": 0.90,
+    })
+    return {"prompt": prompt, "completion": completion}
+
+
+def prepare_training_examples(data: list[dict]) -> list[dict[str, str]]:
+    """Deduplicate sanitized prompts and reject contradictory completions."""
+    grouped: dict[str, dict[str, dict[str, str]]] = {}
+    for item in data:
+        example = format_training_example(item)
+        grouped.setdefault(example["prompt"], {})[example["completion"]] = example
+    return [
+        next(iter(grouped[prompt].values()))
+        for prompt in sorted(grouped)
+        if len(grouped[prompt]) == 1
+    ]
+
+
+def balance_training_examples(
+    examples: list[dict[str, str]],
+    *,
+    income_target: int = 20,
+) -> list[dict[str, str]]:
+    """Oversample represented income classes after conflict-safe deduplication."""
+    by_category: dict[str, list[dict[str, str]]] = {}
+    for example in examples:
+        category = json.loads(example["completion"])["category"]
+        by_category.setdefault(category, []).append(example)
+
+    balanced = list(examples)
+    for category in sorted(INCOME_CATEGORIES):
+        category_examples = by_category.get(category, [])
+        if not category_examples or len(category_examples) >= income_target:
+            continue
+        balanced.extend(
+            category_examples[index % len(category_examples)]
+            for index in range(income_target - len(category_examples))
+        )
+    return balanced
+
+
+def load_training_data():
+    """Load, sanitize, deduplicate, balance, and format labeled transactions."""
+    from datasets import Dataset
+
+    with open(DATA_PATH, encoding="utf-8") as handle:
+        data = json.load(handle)
+    return Dataset.from_list(balance_training_examples(prepare_training_examples(data)))
+
+
+def _load_trainable_model(*, fresh: bool):
+    import torch
+    from peft import LoraConfig, PeftModel, TaskType, get_peft_model
+    from transformers import AutoModelForCausalLM
+
+    base_model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16,
+        device_map="mps",
+        trust_remote_code=True,
+    )
+    adapter_file = OUTPUT_DIR / "adapter_model.safetensors"
+    if adapter_file.exists() and not fresh:
+        print(f"Continuing from adapter: {OUTPUT_DIR}")
+        return PeftModel.from_pretrained(base_model, str(OUTPUT_DIR), is_trainable=True)
+
+    print("Starting a fresh LoRA adapter")
+    return get_peft_model(
+        base_model,
+        LoraConfig(
+            task_type=TaskType.CAUSAL_LM,
+            r=8,
+            lora_alpha=16,
+            lora_dropout=0.05,
+            bias="none",
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+        ),
+    )
+
+
+def main(*, epochs: float = 2.0, fresh: bool = False) -> None:
+    from transformers import AutoTokenizer
+    from trl import SFTConfig, SFTTrainer
+
+    print(f"Loading model: {MODEL_NAME}")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    tokenizer.pad_token = tokenizer.eos_token
+    model = _load_trainable_model(fresh=fresh)
+    model.print_trainable_parameters()
+
+    print("Loading training data...")
+    dataset = load_training_data()
+    company_labels = sum(
+        json.loads(completion)["company_name"] is not None
+        for completion in dataset["completion"]
+    )
+    print(f"Training samples: {len(dataset)}; company labels: {company_labels}")
+
+    trainer = SFTTrainer(
+        model=model,
+        args=SFTConfig(
+            output_dir=str(OUTPUT_DIR),
+            num_train_epochs=epochs,
+            per_device_train_batch_size=2,
+            gradient_accumulation_steps=8,
+            learning_rate=1e-4 if not fresh else 2e-4,
+            warmup_ratio=0.05,
+            logging_steps=10,
+            save_strategy="epoch",
+            save_total_limit=2,
+            bf16=False,
+            fp16=False,
+            optim="adamw_torch",
+            report_to="none",
+            max_length=512,
+        ),
+        train_dataset=dataset,
+        processing_class=tokenizer,
+    )
+
+    print("Starting continued training..." if not fresh else "Starting training...")
+    trainer.train()
+    print(f"Saving LoRA adapter to {OUTPUT_DIR}")
+    model.save_pretrained(str(OUTPUT_DIR))
+    tokenizer.save_pretrained(str(OUTPUT_DIR))
+    print("Done! LoRA adapter saved.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--epochs", type=float, default=2.0)
+    parser.add_argument("--fresh", action="store_true")
+    arguments = parser.parse_args()
+    main(epochs=arguments.epochs, fresh=arguments.fresh)