Update app/main.py

app/main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
@@ -6,49 +6,40 @@ import torch
 
 app = FastAPI()
 
-# ✅ Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2", use_auth_token=True)
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 
-# ✅ Load base model without quantization (for CPU)
 model = AutoModelForCausalLM.from_pretrained(
-    "mistralai/Mistral-7B-Instruct-v0.2",
+    "Qwen/Qwen2.5-0.5B-Instruct",
     torch_dtype=torch.float32,
-    use_auth_token=True
+    trust_remote_code=True
 )
-
-# ✅ Load LoRA adapter
-ADAPTER_DIR = "./adapter/version 1"
-model = PeftModel.from_pretrained(model, ADAPTER_DIR)
+model = PeftModel.from_pretrained(model, "./adapter", is_trainable=False)
 model.eval()
 
-# ✅ Build prompt from messages
 def build_prompt(messages):
     prompt = ""
     for msg in messages:
-        if msg["role"] == "user":
-            prompt += f"### User:\n{msg['content']}\n"
-        elif msg["role"] == "assistant":
-            prompt += f"### Assistant:\n{msg['content']}\n"
+        role = "User" if msg["role"] == "user" else "Assistant"
+        prompt += f"### {role}:\n{msg['content']}\n"
     prompt += "### Assistant:\n"
     return prompt
 
-# ✅ Input format
 class ChatRequest(BaseModel):
-    messages: list  # list of {"role": "user"/"assistant", "content": "..."}
+    messages: list  # [{"role": "user", "content": "..."}]
 
 @app.post("/chat")
 async def chat(req: ChatRequest):
     prompt = build_prompt(req.messages)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    output = model.generate(
+    outputs = model.generate(
         **inputs,
         max_new_tokens=256,
         do_sample=True,
         temperature=0.7,
         top_p=0.95,
-        eos_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id
     )
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    reply = response.split("### Assistant:")[-1].strip()
+    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    reply = output_text.split("### Assistant:")[-1].strip()
     return {"response": reply}