Update app.py

app.py

@@ -1,8 +1,8 @@
+import torch
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from transformers import pipeline
-import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 main = FastAPI()
 
@@ -17,12 +17,12 @@ main.add_middleware(
 MODEL_ID = "tencent/Hy-MT1.5-1.8B-2bit"
 print(f"🔱 Specialist, Loading {MODEL_ID} on CPU...")
 
-# 🔱 Pipeline එක නිවැරදිව load කිරීම
-pipe = pipeline(
-    "text-generation",
-    model=MODEL_ID,
-    device_map="cpu",
-    trust_remote_code=True # කෙලින්ම මෙතනට පමණක් ලබා දෙන්න
+# Load Model & Tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID, 
+    device_map="cpu", 
+    trust_remote_code=True
 )
 
 class ChatRequest(BaseModel):
@@ -32,28 +32,30 @@ class ChatRequest(BaseModel):
 async def chat(request_data: ChatRequest):
     user_query = request_data.message.strip()
     
-    # Simple Prompt
+    # පිරිසිදු Prompt එකක්
     prompt = f"User: {user_query}\nAssistant:"
+    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+    # Output එකෙන් prompt එක අයින් කර පිළිතුර පමණක් ගැනීම
+    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    reply = full_response.split("Assistant:")[-1].strip()
+
+    # පිළිතුර හිස් නම් fallback එකක්
+    if not reply:
+        reply = "I am processing your request, MINZO-PRIME. Please ask again."
 
-    # Response Generation
-    results = pipe(
-        prompt, 
-        max_new_tokens=150, 
-        do_sample=True, 
-        temperature=0.7,
-        pad_token_id=50256
-    )
-
-    # Result Cleaning
-    generated_text = results[0]['generated_text']
-    # Assistant: කියන කොටසෙන් පස්සේ තියෙන ටික විතරක් ගැනීම
-    if "Assistant:" in generated_text:
-        reply = generated_text.split("Assistant:")[-1].strip()
-    else:
-        reply = generated_text.replace(prompt, "").strip()
-
-    return {"reply": reply}
+    print(f"🔱 Generated Reply: {reply}")
+    return {"reply": str(reply)}
 
 @main.get("/")
-def health():
+def home():
     return {"status": "Inachi-Lite Online", "model": MODEL_ID}