Update app.py

app.py

@@ -5,12 +5,12 @@ import torch
 import os
 import json
 import datetime
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from duckduckgo_search import DDGS
 
 app = FastAPI()
 
-# CORS Settings
+# CORS Fix
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -18,31 +18,21 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- දත්ත ගබඩාව සහ මතකය (Learning Path) ---
-API_KEYS_DB = {
-    "ELE-PRIME-ADMIN-SYS": {"credits": 999999, "status": "active"}
-}
+# --- Neural Vault Setup ---
+API_KEYS_DB = {"ELE-PRIME-ADMIN-SYS": {"credits": 999999, "status": "active"}}
 ADMIN_SECRET = "MINZO-SECRET-2026"
-LEARNING_VAULT_PATH = "neural_learning_data.jsonl" # මෙතන තමයි AI එක ඉගෙන ගන්න දත්ත Save වෙන්නේ
+LEARNING_VAULT_PATH = "neural_learning_data.jsonl"
 
-# --- AI Model සැකසුම් ---
-model_id = "mistralai/Mistral-7B-v0.3"
-HF_TOKEN = os.getenv("HF_TOKEN")
+# --- New Model Setup (CPU Optimized) ---
+# Qwen 1.5B එක Mistral වලට වඩා 5 ගුණයකට වඩා සැහැල්ලුයි
+model_id = "Qwen/Qwen2.5-1.5B-Instruct" 
 
-quant_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True,
-)
-
-print("🐘 Elephant Learning Engine Loading...")
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
+print("🐘 Elephant Light-Engine Loading on CPU...")
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    quantization_config=quant_config,
-    device_map="auto",
-    token=HF_TOKEN
+    torch_dtype="auto", # CPU එකට ගැලපෙන පරිදි auto-select වේ
+    device_map="cpu"    # GPU නැති නිසා අනිවාර්යයෙන් CPU ලෙස දිය යුතුයි
 )
 
 # --- Helpers ---
@@ -57,72 +47,52 @@ def get_live_data(query):
             return "\n".join(results)
     except: return ""
 
-# 🧠 CONTINUOUS LEARNING FUNCTION
-# පද්ධතිය විසින් අලුත් දැනුම ගබඩා කරගන්නා ආකාරය
 def capture_learning_data(query, context, response, key_id):
-    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    learning_entry = {
-        "timestamp": timestamp,
-        "key_node": key_id,
-        "instruction": query,
-        "external_context": context,
-        "trained_output": response
+    entry = {
+        "timestamp": str(datetime.datetime.now()),
+        "key": key_id,
+        "q": query,
+        "ctx": context,
+        "ans": response
     }
-    # JSONL ගොනුවකට දත්ත එකතු කිරීම
-    with open(LEARNING_VAULT_PATH, "a", encoding="utf-8") as f:
-        f.write(json.dumps(learning_entry) + "\n")
-    print(f"📊 Neural Data Captured via {key_id}")
+    with open(LEARNING_VAULT_PATH, "a") as f:
+        f.write(json.dumps(entry) + "\n")
 
 # --- Endpoints ---
 
 @app.get("/")
-def home():
-    # පද්ධතිය කොච්චර ඉගෙන ගෙන තියෙනවද කියලා බලාගන්න
-    learning_count = 0
-    if os.path.exists(LEARNING_VAULT_PATH):
-        with open(LEARNING_VAULT_PATH, "r") as f:
-            learning_count = sum(1 for line in f)
-            
-    return {
-        "status": "Elephant AI Node 2026 Live",
-        "active_keys": len(API_KEYS_DB),
-        "learning_entries_captured": learning_count
-    }
+def status():
+    return {"status": "Elephant-Qwen Node Active", "engine": "Qwen-2.5-1.5B"}
 
 @app.post("/admin/add-key")
-async def register_key(data: NewKeyRequest):
-    if data.admin_pass != ADMIN_SECRET:
-        raise HTTPException(status_code=401, detail="Unauthorized Admin Access")
+async def add_key(data: NewKeyRequest):
+    if data.admin_pass != ADMIN_SECRET: raise HTTPException(status_code=401)
     API_KEYS_DB[data.new_key] = {"credits": 5000, "status": "active"}
-    return {"message": f"Key {data.new_key} Registered Successfully"}
+    return {"message": "Key Registered"}
 
 @app.post("/v1/chat")
-async def chat_endpoint(message: dict, x_api_key: str = Header(None)):
-    if x_api_key not in API_KEYS_DB:
-        raise HTTPException(status_code=403, detail="Invalid API Key")
+async def chat(message: dict, x_api_key: str = Header(None)):
+    if x_api_key not in API_KEYS_DB: raise HTTPException(status_code=403)
     
-    user_query = message.get("query", "")
-    context = ""
-    
-    # 🌐 2026 Live Web Search
-    if any(word in user_query.lower() for word in ["today", "now", "2026", "news", "price"]):
-        context = get_live_data(user_query)
+    query = message.get("query", "")
+    context = get_live_data(query) if any(w in query.lower() for w in ["now", "2026", "today"]) else ""
+
+    # Qwen Chat Template එක භාවිතය
+    messages = [
+        {"role": "system", "content": f"You are Elephant AI. Current year: 2026. Context: {context}"},
+        {"role": "user", "content": query}
+    ]
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    model_inputs = tokenizer([text], return_tensors="pt").to("cpu")
 
-    full_prompt = f"System: Year 2026. Context: {context}\nUser: {user_query}\nAssistant:"
-    inputs = tokenizer(full_prompt, return_tensors="pt").to("cuda")
-    
     with torch.no_grad():
-        output_tokens = model.generate(**inputs, max_new_tokens=300, do_sample=True, temperature=0.7)
+        generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=512)
+        response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
     
-    response = tokenizer.decode(output_tokens[0], skip_special_tokens=True).split("Assistant:")[-1].strip()
-
-    # 🧠 දියුණු වීමට අවශ්‍ය දත්ත ගබඩා කිරීම (Learning Trigger)
-    capture_learning_data(user_query, context, response, x_api_key)
-
-    return {
-        "reply": response, 
-        "key_id": x_api_key,
-        "learning_status": "synced"
-    }
+    # පද්ධතියේ පිළිතුර පිරිසිදු කිරීම
+    clean_response = response.split("assistant")[-1].strip()
+    
+    capture_learning_data(query, context, clean_response, x_api_key)
+    return {"reply": clean_response, "key_id": x_api_key}
 
 main = app