Update fine_tune.py

fine_tune.py

@@ -1,6 +1,5 @@
 # fine_tune.py
 import os
-os.environ["OMP_NUM_THREADS"] = "1"
 import torch
 from datasets import load_dataset
 from transformers import (
@@ -14,19 +13,25 @@ from trl import SFTTrainer
 from fastapi import FastAPI
 import uvicorn
 
+
 base_model_name = "unsloth/llama-3-8b-Instruct-bnb-4bit"
-output_dir = "/data/fine_tuning"   # persist inside /code
 
+output_dir = "/data/fine_tuning" 
 dataset_path = "dataset.jsonl"
 
-# --- Check if already fine-tuned ---
-if not os.path.exists(output_dir):
+# Initialize model and tokenizer variables
+model = None
+tokenizer = None
+
+# 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...")
 
     # Load dataset
     dataset = load_dataset("json", data_files=dataset_path, split="train")
 
-    # Load base model + tokenizer
+    # Load base model for training
     model = AutoModelForCausalLM.from_pretrained(
         base_model_name,
         device_map="auto",
@@ -41,9 +46,9 @@ if not os.path.exists(output_dir):
 
     # Configure LoRA
     peft_config = LoraConfig(
+        r=16,
         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"],
@@ -58,7 +63,6 @@ if not os.path.exists(output_dir):
         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,
@@ -67,30 +71,46 @@ if not os.path.exists(output_dir):
         lr_scheduler_type="linear",
     )
 
-    # Train
+    # Initialize Trainer
     trainer = SFTTrainer(
         model=model,
         train_dataset=dataset,
         peft_config=peft_config,
+        dataset_text_field="text", # Ensure your dataset has a 'text' column
         args=training_arguments,
     )
+    
+    # Train the model
     trainer.train()
+    
+    # Save the trained adapter
     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 = PeftModel.from_pretrained(base_model, output_dir)
-tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
+    print(f"Fine-tuned model adapter saved to {output_dir}")
+    
+    model = trainer.model
+
+# Inference Logic
+# If training did not run, load the existing model
+else:
+    print("Found existing fine-tuned model. Loading for inference...")
+    
+    # Load the base model
+    base_model = AutoModelForCausalLM.from_pretrained(
+        base_model_name,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+    # Apply the PEFT adapter
+    model = PeftModel.from_pretrained(base_model, output_dir)
+    tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
+
 
+# 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
 app = FastAPI(title="Fine-tuned LLaMA API")
 
 @app.get("/")
@@ -99,8 +119,10 @@ def home():
 
 @app.post("/generate")
 def generate(prompt: str):
-    outputs = pipe(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
+    
+    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"]}
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)