Update app.py

app.py

@@ -1,14 +1,12 @@
-import torch
-from fastapi import FastAPI, Header, HTTPException
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import datetime
+from transformers import pipeline
+import torch
 
-# ── API INITIALIZATION ──
-# Hugging Face සොයන 'main' attribute එක මෙතනට ලබා දී ඇත
 main = FastAPI()
 
+# CORS සක්‍රීය කිරීම - ඕනෑම තැනක සිට සම්බන්ධ වීමට
 main.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -16,71 +14,45 @@ main.add_middleware(
     allow_headers=["*"],
 )
 
-# ── CONFIGURATION ──
-API_KEYS_DB = {
-    "ELE-PRIME-ADMIN-SYS": {"status": "active"},
-}
 MODEL_ID = "tencent/Hy-MT1.5-1.8B-2bit"
-
-print(f"🔱 Specialist, Loading {MODEL_ID} on CPU...")
-
-# Load Tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, clean_up_tokenization_spaces=False) 
-                                          
-# Load Model Optimized for CPU
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    device_map="cpu", # GPU නොමැති නිසා කෙලින්ම CPU එකට ලබා දීම
-    torch_dtype=torch.float32, 
-    trust_remote_code=True
+print(f"🔱 Specialist, Loading {MODEL_ID} (Public Mode)...")
+
+# Pipeline initialization
+pipe = pipeline(
+    "text-generation",
+    model=MODEL_ID,
+    device_map="cpu",
+    model_kwargs={"trust_remote_code": True}
 )
 
-print("🔱 Inachi-Lite is Online and Ready.")
-
-# ── DATA MODELS ──
 class ChatRequest(BaseModel):
     message: str
-    history: list = []
-    temperature: float = 0.7
-    max_tokens: int = 512
 
-# ── CHAT ENDPOINT ──
 @main.post("/v1/chat")
-async def chat(request_data: ChatRequest, x_api_key: str = Header(None)):
-    # API Key පරීක්ෂා කිරීම
-    if not x_api_key or x_api_key not in API_KEYS_DB:
-        raise HTTPException(status_code=403, detail="Access Denied")
-
+async def chat(request_data: ChatRequest):
     user_query = request_data.message.strip()
-    today = datetime.datetime.now().strftime("%Y-%m-%d")
-
-    # Prompt සකස් කිරීම
-    prompt = f"System: You are Inachi AI, an expert assistant for MINZO-PRIME. Date: {today}\n"
-    
-    # History ඇතුළත් කිරීම
-    for human, ai in request_data.history[-2:]: 
-        prompt += f"User: {human}\nAI: {ai}\n"
     
-    prompt += f"User: {user_query}\nAI:"
+    # Prompt Setup
+    prompt = f"User: {user_query}\nAssistant:"
 
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
+    # Text Generation
+    results = pipe(
+        prompt, 
+        max_new_tokens=150, 
+        do_sample=True, 
+        temperature=0.7,
+        pad_token_id=50256
+    )
 
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=request_data.max_tokens,
-            temperature=request_data.temperature,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
+    # Result Cleaning
+    generated_text = results[0]['generated_text']
+    reply = generated_text.split("Assistant:")[-1].strip()
 
-    response_text = tokenizer.decode(outputs[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True).strip()
-    
-    return {
-        "reply": response_text,
-        "status": "success"
-    }
+    if not reply:
+        reply = "I am here, MINZO-PRIME. Systems are nominal."
+
+    return {"reply": reply}
 
 @main.get("/")
-def home():
-    return {"message": "Inachi-Lite is Running", "status": "online"}
+def health():
+    return {"status": "Public Inachi-Lite Online"}