Update fine_tune.py

fine_tune.py

@@ -1,6 +1,4 @@
-
-# fine_tuning/fine_tune.py
-
+# fine_tune.py
 import os
 import torch
 from datasets import load_dataset
@@ -8,85 +6,100 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     TrainingArguments,
-    BitsAndBytesConfig
+    pipeline
 )
-from peft import LoraConfig
+from peft import LoraConfig, PeftModel
 from trl import SFTTrainer
-
+from fastapi import FastAPI
+import uvicorn
 
 base_model_name = "unsloth/llama-3-8b-Instruct-bnb-4bit"
-output_dir = "/tmp/fine_tuning"
+output_dir = "fine_tuning"   # persist inside /code
 
 dataset_path = "dataset.jsonl"
 
-
-# Load the Dataset
-print("Loading dataset...")
-dataset = load_dataset("json", data_files=dataset_path, split="train")
-print("Dataset loaded successfully.")
-
-
-# Load the Base Model & Tokenizer
-print(f"Loading base model: {base_model_name}...")
-model = AutoModelForCausalLM.from_pretrained(
+# --- Check if already fine-tuned ---
+if not os.path.exists(output_dir):
+    print("No fine-tuned model found. Starting training...")
+
+    # Load dataset
+    dataset = load_dataset("json", data_files=dataset_path, split="train")
+
+    # Load base model + tokenizer
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model_name,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+    model.config.use_cache = False
+
+    tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "right"
+
+    # Configure LoRA
+    peft_config = LoraConfig(
+        lora_alpha=16,
+        lora_dropout=0.1,
+        r=16,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+    )
+
+    # Training args
+    training_arguments = TrainingArguments(
+        output_dir=output_dir,
+        num_train_epochs=1,
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=2,
+        optim="paged_adamw_32bit",
+        logging_steps=10,
+        learning_rate=2e-4,
+        weight_decay=0.01,
+        fp16=True,
+        max_grad_norm=0.3,
+        max_steps=-1,
+        warmup_ratio=0.03,
+        group_by_length=True,
+        lr_scheduler_type="linear",
+    )
+
+    # Train
+    trainer = SFTTrainer(
+        model=model,
+        train_dataset=dataset,
+        peft_config=peft_config,
+        args=training_arguments,
+    )
+    trainer.train()
+    trainer.model.save_pretrained(output_dir)
+    print(f"Fine-tuned model saved to {output_dir}")
+
+# --- Load model for inference ---
+print("Loading fine-tuned model...")
+base_model = AutoModelForCausalLM.from_pretrained(
     base_model_name,
     device_map="auto",
     trust_remote_code=True,
 )
-model.config.use_cache = False
-print("Base model loaded successfully.")
-
+model = PeftModel.from_pretrained(base_model, output_dir)
 tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-tokenizer.padding_side = "right"
-print("Tokenizer loaded and configured.")
 
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
 
-# Configure PEFT (LoRA)
-peft_config = LoraConfig(
-    lora_alpha=16,
-    lora_dropout=0.1,
-    r=16,
-    bias="none",
-    task_type="CAUSAL_LM",
-    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
-)
-print("LoRA configured.")
-
-
-# Define Training Arguments
-training_arguments = TrainingArguments(
-    output_dir=output_dir,
-    num_train_epochs=1,
-    per_device_train_batch_size=2,
-    gradient_accumulation_steps=2,
-    optim="paged_adamw_32bit",
-    logging_steps=10,
-    learning_rate=2e-4,
-    weight_decay=0.01,
-    fp16=True,
-    max_grad_norm=0.3,
-    max_steps=-1,
-    warmup_ratio=0.03,
-    group_by_length=True,
-    lr_scheduler_type="linear",
-)
-print("Training arguments set.")
-
+# --- FastAPI app ---
+app = FastAPI(title="Fine-tuned LLaMA API")
 
-# Initialize and Start Training
-trainer = SFTTrainer(
-    model=model,
-    train_dataset=dataset,
-    peft_config=peft_config,
-    args=training_arguments,
-)
-print("Trainer initialized. Starting the fine-tuning process...")
-trainer.train()
-print("Training complete.")
+@app.get("/")
+def home():
+    return {"status": "ok", "message": "Fine-tuned LLaMA is ready."}
 
+@app.post("/generate")
+def generate(prompt: str):
+    outputs = pipe(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
+    return {"response": outputs[0]["generated_text"]}
 
-# Save the Final Model
-trainer.model.save_pretrained(output_dir)
-print(f"Fine-tuned model adapter saved to {output_dir}")
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)