Update main.py

main.py

@@ -1,26 +1,24 @@
-import os
-os.environ["HF_HOME"] = "/tmp"
-
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel
-from fastapi import FastAPI
-from pydantic import BaseModel
-
-model_name = "microsoft/phi-2"
-adapter_path = "howtomakepplragequit/phi2-lora-instruct"
-
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-base_model = AutoModelForCausalLM.from_pretrained(model_name)
-model = PeftModel.from_pretrained(base_model, adapter_path)
-
-app = FastAPI()
-
-class Prompt(BaseModel):
-    input: str
-
-@app.post("/chat")
-def chat(prompt: Prompt):
-    inputs = tokenizer(prompt.input, return_tensors="pt")
-    output = model.generate(**inputs, max_new_tokens=50)
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    return {"response": response}
+from fastapi import FastAPI, Request
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel, PeftConfig
+import torch
+
+app = FastAPI()
+
+model_name = "microsoft/phi-2"
+peft_model_id = "howtomakepplragequit/phi2-lora-instruct"
+
+# Load tokenizer and model with LoRA
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+base_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+model = PeftModel.from_pretrained(base_model, peft_model_id)
+model.eval()
+
+@app.post("/generate")
+async def generate(request: Request):
+    data = await request.json()
+    prompt = data.get("prompt", "")
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=100)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"response": response}