Update fine_tune.py

fine_tune.py

@@ -105,35 +105,88 @@ print("Inference pipeline ready.")
 class GenerateRequest(BaseModel):
     prompt: str
 
-app = FastAPI(title="Fine-tuned LLaMA API")
+app = FastAPI(
+    title="DirectEd AI Assistant",
+    version="1.0",
+    description="API for fine-tuned DirectEd AI chatbot."
+)
+
+# --- Load Model + Tokenizer ---
+try:
+    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model_name,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto"
+    )
 
+    if os.path.exists(output_dir):
+        print(f"Loading adapter from {output_dir}")
+        model = PeftModel.from_pretrained(model, output_dir)
+    else:
+        print("⚠️ No adapter folder found, using base model only")
+
+except Exception as e:
+    print("❌ Model load failed:", e)
+    model, tokenizer = None, None
+
+
+# --- Routes ---
 @app.get("/")
-def home():
-    return {"status": "ok", "message": "Fine-tuned LLaMA is ready."}
+def health():
+    return {"status": "ok", "message": "DirectEd AI Space running."}
+
 
 @app.post("/generate")
-def generate(request: GenerateRequest):
-    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)
+def generate(prompt: str, max_new_tokens: int = 200):
+    if model is None or tokenizer is None:
+        return {"error": "Model not loaded."}
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        top_k=50,
+        top_p=0.9
+    )
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"response": text}
+
+
+@app.get("/list_adapter")
+def list_adapter():
+    """List adapter files in output_dir"""
+    if os.path.exists(output_dir):
+        files = os.listdir(output_dir)
+        return {"adapter_files": files}
+    return {"adapter_files": [], "message": "No adapter directory found."}
+
+
+@app.post("/upload_adapter")
+def upload_adapter(file: UploadFile = File(...)):
+    """Upload adapter files (e.g. adapter_config.json, adapter_model.bin)"""
+    os.makedirs(output_dir, exist_ok=True)
+    save_path = os.path.join(output_dir, file.filename)
+    with open(save_path, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
+    return {"status": "success", "filename": file.filename}
+
+
+@app.post("/push_adapter")
+def push_adapter():
+    """Push adapter folder to Hugging Face Hub"""
+    if not os.path.exists(output_dir):
+        return {"error": "No adapter folder found."}
+
+    files = os.listdir(output_dir)
+    if not files:
+        return {"error": "Adapter folder is empty."}
+
+    upload_folder(
+        repo_id=hub_repo_id,
+        folder_path=output_dir,
+        commit_message="Upload LoRA adapter from Space"
+    )
+    return {"status": "uploaded", "repo": f"https://huggingface.co/{hub_repo_id}", "files": files}