Create app/train.py

backend/app/train.py

@@ -0,0 +1,79 @@
+import json
+import os
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling,
+)
+import torch
+from datasets import Dataset
+
+# -------- Settings --------
+MODEL_NAME = "google/gemma-1.1-2b-it"
+DATA_PATH = "./backend/data/pregnancy_dataset.json"
+SAVE_PATH = "./backend/app/checkpoints"
+
+# -------- Load Dataset --------
+def load_dataset():
+    with open(DATA_PATH, "r") as f:
+        data = json.load(f)
+    dataset = Dataset.from_list(data)
+    return dataset
+
+# -------- Tokenization --------
+def tokenize_function(example, tokenizer):
+    return tokenizer(
+        f"<start_of_turn> {example['prompt']} <end_of_turn>\n{example['completion']} <end_of_turn>",
+        truncation=True,
+        padding="max_length",
+        max_length=256,
+    )
+
+# -------- Main Training Function --------
+def train():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+
+    # Optional: Freeze all except final layers
+    for param in model.base_model.parameters():
+        param.requires_grad = False
+    for param in model.lm_head.parameters():
+        param.requires_grad = True
+
+    # Load and tokenize
+    raw_dataset = load_dataset()
+    tokenized_dataset = raw_dataset.map(lambda x: tokenize_function(x, tokenizer))
+
+    # Data collator
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir=SAVE_PATH,
+        num_train_epochs=3,
+        per_device_train_batch_size=2,
+        save_steps=50,
+        logging_steps=10,
+        save_total_limit=1,
+        remove_unused_columns=False,
+        report_to="none",
+    )
+
+    # Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+    )
+
+    trainer.train()
+    trainer.save_model(SAVE_PATH)
+    tokenizer.save_pretrained(SAVE_PATH)
+    print("✅ Fine-tuned model saved!")
+
+if __name__ == "__main__":
+    train()