Update fine_tune.py

fine_tune.py

@@ -9,9 +9,6 @@ from transformers import (
 )
 from peft import LoraConfig, PeftModel
 from trl import SFTTrainer
-from fastapi import FastAPI, UploadFile, File
-from huggingface_hub import upload_folder
-import shutil
 from fastapi import FastAPI
 from pydantic import BaseModel
 import uvicorn
@@ -108,88 +105,14 @@ print("Inference pipeline ready.")
 class GenerateRequest(BaseModel):
     prompt: str
 
-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"
-    )
+app = FastAPI(title="Fine-tuned LLaMA API")
 
-    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 health():
-    return {"status": "ok", "message": "DirectEd AI Space running."}
-
+def home():
+    return {"status": "ok", "message": "Fine-tuned LLaMA is ready."}
 
 @app.post("/generate")
-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}
+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"]}