Update app.py

app.py

@@ -3,12 +3,9 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import torch
 import os
-import json
 import re
-import uuid
 import secrets
-import datetime
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from duckduckgo_search import DDGS
 
 app = FastAPI()
@@ -20,100 +17,60 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- Database & Config ---
-# ආරම්භක Keys
+# --- Specialist DB ---
 API_KEYS_DB = {
-    "ELE-PRIME-ADMIN-SYS": {"limit": 10000, "used": 0, "status": "active"},
-    "ELE-PRIME-YG5EPZFQ": {"limit": 5000, "used": 0, "status": "active"}
+    "ELE-PRIME-ADMIN-SYS": {"limit": 100000, "used": 0, "status": "active"},
+    "ELE-PRIME-VOID-X": {"limit": 50000, "used": 0, "status": "active"}
 }
 ADMIN_SECRET = "MINZO-SECRET-2026"
 
-# --- AI Model ---
-model_id = "google/gemma-3-1b-it"
-print(f"🔱 INACHI-CORE: Loading {model_id}...")
+# --- MiMo-Audio 7B Optimization ---
+model_id = "XiaomiMiMo/MiMo-Audio-7B-Instruct"
+print(f"🔱 INACHI-CORE: Deploying Multimodal Engine {model_id}...")
+
+# 4-bit Quantization එකතු කිරීම (CPU/RAM ඉතිරි කිරීමට)
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16
+)
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="cpu")
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    quantization_config=quant_config,
+    device_map="auto" # මෙය ස්වයංක්‍රීයව RAM එක කළමනාකරණය කරයි
+)
 
-# --- Data Models ---
 class AdminRequest(BaseModel):
     admin_pass: str
-    limit: int = 1000
-
-# --- API Endpoints ---
-
-@app.get("/")
-def home():
-    return {"status": "Elephant Pro Active", "active_keys": len(API_KEYS_DB)}
-
-# 🔱 අලුතින් Key එකක් Auto-Generate කරන Endpoint එක
-@app.post("/v1/generate-key")
-async def generate_key(data: AdminRequest):
-    if data.admin_pass != ADMIN_SECRET:
-        raise HTTPException(status_code=401, detail="Unauthorized Specialist Access!")
-    
-    # Random Key එකක් නිර්මාණය කිරීම (උදා: ELE-PRIME-X8A2...)
-    new_key = f"ELE-PRIME-{secrets.token_hex(4).upper()}"
-    API_KEYS_DB[new_key] = {"limit": data.limit, "used": 0, "status": "active"}
-    
-    return {
-        "message": "New Specialist Key Activated",
-        "api_key": new_key,
-        "limit": data.limit
-    }
+    limit: int = 5000
 
 @app.post("/v1/chat")
 async def chat(message: dict, x_api_key: str = Header(None)):
     if not x_api_key or x_api_key not in API_KEYS_DB:
-        raise HTTPException(status_code=403, detail="Access Denied")
+        raise HTTPException(status_code=403, detail="Invalid Key")
         
-    key_info = API_KEYS_DB[x_api_key]
-    if key_info["used"] >= key_info["limit"]:
-        raise HTTPException(status_code=429, detail="Limit Reached")
-    
     query = message.get("query", "")
-
-    # Web Search
-    context = ""
-    if any(w in query.lower() for w in ["today", "now", "2026", "අද"]):
-        try:
-            with DDGS() as ddgs:
-                results = list(ddgs.text(query, max_results=2))
-                context = "\n".join([r['body'] for r in results])
-        except: pass
-
-    # 🔱 Language Adaptive System Instruction
+    
+    # MiMo පද්ධතියේ System Prompt එක
     system_instruction = (
-        "You are Elephant AI (Inachi-Core), an expert assistant for Specialist MINZO-PRIME. "
-        "Respond in the language used by the user (Sinhala or English). "
-        f"Real-time Context: {context}"
+        "You are Inachi-Prime, an Any-to-Any multimodal AI assistant. "
+        "You are designed by Specialist MINZO-PRIME. "
+        "Respond accurately in Sinhala or English based on user input."
     )
 
-    msgs = [
-        {"role": "system", "content": system_instruction},
-        {"role": "user", "content": query}
-    ]
-    
-    text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer([text], return_tensors="pt").to("cpu")
+    inputs = tokenizer(f"{system_instruction}\n\nUser: {query}\nAssistant:", return_tensors="pt").to("cpu")
 
     with torch.no_grad():
         outputs = model.generate(
-            inputs.input_ids, 
+            **inputs,
             max_new_tokens=512,
-            temperature=0.6, 
-            top_p=0.9,
+            temperature=0.7,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id
         )
         
-        full_response = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-        ans = full_response.split("assistant")[-1].strip()
-        
-        # Cleaning Logic
-        if "</think>" in ans: ans = ans.split("</think>")[-1].strip()
-        ans = ans.replace("Ċ", "\n").replace("Ġ", " ")
-        ans = re.sub(r' +', ' ', ans).strip()
+        ans = tokenizer.decode(outputs[0], skip_special_tokens=True).split("Assistant:")[-1].strip()
 
     API_KEYS_DB[x_api_key]["used"] += 1
     return {"reply": ans, "usage": API_KEYS_DB[x_api_key]["used"]}