Update app/main.py

app/main.py

@@ -1,45 +1,57 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+from peft import PeftModel
+import json
+import os
 
-app = FastAPI()
-
-tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", trust_remote_code=True)
-tokenizer.pad_token = tokenizer.eos_token
-
+# Load tokenizer and base model
+base_model = "Qwen/Qwen2-0.5B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
-    "Qwen/Qwen2.5-0.5B-Instruct",
-    torch_dtype=torch.float32,
+    base_model,
+    device_map="cuda" if torch.cuda.is_available() else "cpu",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
     trust_remote_code=True
 )
+
+# Clean up adapter_config.json before loading adapter
+adapter_config_path = "./adapter/adapter_config.json"
+if os.path.exists(adapter_config_path):
+    with open(adapter_config_path, "r") as f:
+        adapter_config = json.load(f)
+    for key in ["corda_config", "eva_config", "megatron_config"]:
+        adapter_config.pop(key, None)
+    with open(adapter_config_path, "w") as f:
+        json.dump(adapter_config, f)
+
+# Load adapter
 model = PeftModel.from_pretrained(model, "./adapter", is_trainable=False)
 model.eval()
 
-def build_prompt(messages):
-    prompt = ""
-    for msg in messages:
-        role = "User" if msg["role"] == "user" else "Assistant"
-        prompt += f"### {role}:\n{msg['content']}\n"
-    prompt += "### Assistant:\n"
-    return prompt
+# Simple chat function
+def chat(prompt):
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": prompt}
+    ]
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-class ChatRequest(BaseModel):
-    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)
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=256,
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=512,
         do_sample=True,
         temperature=0.7,
-        top_p=0.95,
-        eos_token_id=tokenizer.eos_token_id
+        streamer=streamer
     )
-    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    reply = output_text.split("### Assistant:")[-1].strip()
-    return {"response": reply}
+    output = tokenizer.decode(generated_ids[0][model_inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    return output
+
+# Example
+if __name__ == "__main__":
+    while True:
+        prompt = input("User: ")
+        if prompt.lower() in ["exit", "quit"]:
+            break
+        print("AI:", chat(prompt))