Rename trainer.py to src/train.py

src/train.py

@@ -0,0 +1,88 @@
+import os
+from typing import Optional
+
+import torch
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    Trainer,
+    TrainingArguments,
+)
+
+from peft import LoraConfig, TaskType, get_peft_model
+
+
+def finetune_lora(
+    base_model: str,
+    dataset_id: str,
+    text_column: str,
+    output_dir: str,
+    max_train_samples: int = 2000,
+    max_steps: int = 100,
+    learning_rate: float = 2e-4,
+    batch_size: int = 2,
+    lora_r: int = 8,
+    lora_alpha: int = 16,
+    lora_dropout: float = 0.05,
+) -> str:
+    ds = load_dataset(dataset_id, split="train")
+    if text_column not in ds.column_names:
+        return f"ERROR: column '{text_column}' not found. Available: {ds.column_names}"
+
+    if max_train_samples and max_train_samples > 0:
+        ds = ds.select(range(min(len(ds), int(max_train_samples))))
+
+    tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    def tok(batch):
+        return tokenizer(batch[text_column], truncation=True, max_length=256)
+
+    tokenized = ds.map(tok, batched=True, remove_columns=ds.column_names)
+
+    model = AutoModelForCausalLM.from_pretrained(base_model)
+    model.config.pad_token_id = tokenizer.pad_token_id
+
+    # LoRA target modules here are GPT-2-ish defaults.
+    # If you swap to a non-GPT2 architecture, you may need to change target_modules.
+    lora_cfg = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        r=int(lora_r),
+        lora_alpha=int(lora_alpha),
+        lora_dropout=float(lora_dropout),
+        bias="none",
+        target_modules=["c_attn", "c_proj"],
+    )
+    model = get_peft_model(model, lora_cfg)
+
+    collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+    fp16 = torch.cuda.is_available()
+    args = TrainingArguments(
+        output_dir=output_dir,
+        per_device_train_batch_size=int(batch_size),
+        learning_rate=float(learning_rate),
+        max_steps=int(max_steps),
+        logging_steps=10,
+        save_steps=0,
+        report_to=[],
+        fp16=fp16,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=args,
+        train_dataset=tokenized,
+        data_collator=collator,
+    )
+
+    trainer.train()
+
+    adapter_dir = os.path.join(output_dir, "adapter")
+    model.save_pretrained(adapter_dir)
+    tokenizer.save_pretrained(adapter_dir)
+
+    return f"Saved LoRA adapter + tokenizer to {adapter_dir}"

trainer.py

@@ -1,58 +0,0 @@
-from datasets import load_dataset
-from transformers import (
-    AutoTokenizer,
-    AutoModelForSequenceClassification,
-    TrainingArguments,
-    Trainer
-)
-from peft import LoraConfig, get_peft_model
-
-
-def run_finetune(base_model, dataset_path, epochs=3):
-
-    dataset = load_dataset("json", data_files=dataset_path)
-
-    tokenizer = AutoTokenizer.from_pretrained(base_model)
-
-    def tokenize(example):
-        return tokenizer(
-            example["text"],
-            truncation=True,
-            padding="max_length"
-        )
-
-    tokenized = dataset.map(tokenize)
-
-    model = AutoModelForSequenceClassification.from_pretrained(
-        base_model,
-        num_labels=2
-    )
-
-    lora_config = LoraConfig(
-        r=8,
-        lora_alpha=32,
-        target_modules=["query", "value"],
-        lora_dropout=0.05
-    )
-
-    model = get_peft_model(model, lora_config)
-
-    args = TrainingArguments(
-        output_dir="./results",
-        num_train_epochs=epochs,
-        per_device_train_batch_size=4,
-        save_steps=50,
-        logging_steps=10
-    )
-
-    trainer = Trainer(
-        model=model,
-        args=args,
-        train_dataset=tokenized["train"]
-    )
-
-    trainer.train()
-
-    model.save_pretrained("./finetuned")
-
-    return "Training complete. Model saved."