Update fine_tune.py

fine_tune.py

@@ -17,20 +17,18 @@ import uvicorn
 base_model_name = "unsloth/llama-3-8b-Instruct-bnb-4bit"
 output_dir = "/data/fine_tuning"
 dataset_path = "dataset.jsonl"
+hub_model_id = "Nutnell/direct-ed-finetune-job"
 
 # --- 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 for training
     model = AutoModelForCausalLM.from_pretrained(
         base_model_name,
         device_map="auto",
@@ -53,7 +51,6 @@ if not os.path.exists(os.path.join(output_dir, 'adapter_config.json')):
         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,
@@ -69,10 +66,9 @@ 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/direct-ed-finetune-job",
+        hub_model_id=hub_model_id,
     )
 
-    # Initialize Trainer
     trainer = SFTTrainer(
         model=model,
         train_dataset=dataset,
@@ -81,40 +77,31 @@ if not os.path.exists(os.path.join(output_dir, 'adapter_config.json')):
         args=training_arguments,
     )
     
-    # Train the model
     trainer.train()
-    
-    # Save the trained adapter
     trainer.model.save_pretrained(output_dir)
     print(f"Fine-tuned model adapter saved to {output_dir}")
     
+    # Push trained model to Hub
+    trainer.push_to_hub()
+
     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 ---
+# --- Inference Pipeline ---
 print("Setting up inference pipeline...")
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
 print("Inference pipeline ready.")
 
-
 # --- FastAPI App ---
-
-# 2. DEFINE THE PYDANTIC MODEL FOR THE REQUEST BODY
 class GenerateRequest(BaseModel):
     prompt: str
 
@@ -124,14 +111,29 @@ app = FastAPI(title="Fine-tuned LLaMA API")
 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(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"
-    
     outputs = pipe(formatted_prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
     return {"response": outputs[0]["generated_text"]}
 
+# --- Extra utility endpoints ---
+@app.get("/list-files")
+def list_files():
+    files = []
+    for root, _, filenames in os.walk(output_dir):
+        for fname in filenames:
+            files.append(os.path.relpath(os.path.join(root, fname), output_dir))
+    return {"files": files}
+
+@app.post("/push-to-hub")
+def push_to_hub():
+    try:
+        model.push_to_hub(hub_model_id)
+        tokenizer.push_to_hub(hub_model_id)
+        return {"status": "success", "message": f"Pushed to Hugging Face Hub ({hub_model_id})"}
+    except Exception as e:
+        return {"status": "error", "message": str(e)}
+
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)