Upload main.py

main.py

@@ -1,27 +1,26 @@
 import os
+os.environ["HF_HOME"] = "/tmp"
+
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-
-# Force Hugging Face cache to a writable dir
-os.environ["HF_HOME"] = "/data"
 
-model_name = "howtomakepplragequit/phi2-lora-instruct"
+model_name = "microsoft/phi-2"
+adapter_path = "howtomakepplragequit/phi2-lora-instruct"
 
-# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Create pipeline
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+base_model = AutoModelForCausalLM.from_pretrained(model_name)
+model = PeftModel.from_pretrained(base_model, adapter_path)
 
-# FastAPI app setup
 app = FastAPI()
 
 class Prompt(BaseModel):
-    prompt: str
+    input: str
 
-@app.post("/generate")
-def generate_text(data: Prompt):
-    output = generator(data.prompt, max_length=200, do_sample=True)[0]["generated_text"]
-    return {"response": output}
+@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}