Update fine_tune.py

fine_tune.py

@@ -1,4 +1,3 @@
-# fine_tune.py
 import os
 import torch
 from datasets import load_dataset
@@ -11,19 +10,19 @@ from transformers import (
 from peft import LoraConfig, PeftModel
 from trl import SFTTrainer
 from fastapi import FastAPI
+from pydantic import BaseModel # 1. ADD THIS IMPORT
 import uvicorn
 
-
+# --- Configuration ---
 base_model_name = "unsloth/llama-3-8b-Instruct-bnb-4bit"
-
-output_dir = "/data/fine_tuning" 
+output_dir = "/data/fine_tuning"
 dataset_path = "dataset.jsonl"
 
-# Initialize model and tokenizer variables
+# --- Initialize model and tokenizer variables ---
 model = None
 tokenizer = None
 
-# Training Logic
+# --- Training Logic ---
 # Check if a fine-tuned model adapter already exists
 if not os.path.exists(os.path.join(output_dir, 'adapter_config.json')):
     print("No fine-tuned model found. Starting training...")
@@ -70,7 +69,7 @@ if not os.path.exists(os.path.join(output_dir, 'adapter_config.json')):
         group_by_length=True,
         lr_scheduler_type="linear",
         push_to_hub=True,
-        hub_model_id="Nutnell/DirectEd-AI", 
+        hub_model_id="Nutnell/DirectEd-AI",
     )
 
     # Initialize Trainer
@@ -91,7 +90,7 @@ if not os.path.exists(os.path.join(output_dir, 'adapter_config.json')):
     
     model = trainer.model
 
-# Inference Logic
+# --- Inference Logic ---
 # If training did not run, load the existing model
 else:
     print("Found existing fine-tuned model. Loading for inference...")
@@ -107,22 +106,30 @@ else:
     tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
 
 
-# Create Inference Pipeline
+# --- Create Inference Pipeline ---
 print("Setting up inference pipeline...")
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
 print("Inference pipeline ready.")
 
-# FastAPI App
+
+# --- FastAPI App ---
+
+# 2. DEFINE THE PYDANTIC MODEL FOR THE REQUEST BODY
+class GenerateRequest(BaseModel):
+    prompt: str
+
 app = FastAPI(title="Fine-tuned LLaMA API")
 
 @app.get("/")
 def home():
     return {"status": "ok", "message": "Fine-tuned LLaMA is ready."}
 
+# 3. UPDATE THE GENERATE FUNCTION TO USE THE PYDANTIC MODEL
 @app.post("/generate")
-def generate(prompt: str):
+def generate(request: GenerateRequest):
+    # Access the prompt from the request object
+    formatted_prompt = f"<|start_header_id|>user<|end_header_id|>\n\n{request.prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
     
-    formatted_prompt = f"<|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
     outputs = pipe(formatted_prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
     return {"response": outputs[0]["generated_text"]}