Update app.py

app.py

@@ -1,130 +1,90 @@
-from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse, JSONResponse
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from peft import PeftModel
-import threading
-import asyncio
+import os
+import requests
+from pathlib import Path
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from llama_cpp import Llama
 
-# ---------------- CONFIG ----------------
-MODEL_ID = "Neon-AI/Kushina"
-MAX_NEW_TOKENS = 16384
+# ================= CONFIG =================
+MODEL_URL = "https://huggingface.co/Neon-AI/Kushina/resolve/main/model.gguf"
+MODEL_PATH = "model.gguf"
+N_CTX = 16384
+N_THREADS = 4
+N_BATCH = 256
+MAX_TOKENS = 16384
 TEMPERATURE = 0.7
 TOP_P = 0.9
+# ==========================================
 
-# ---------------- LOAD MODEL ----------------
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-base_model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32, device_map=None)
+app = FastAPI(title="Kushina API", version="1.0")
 
-# Try loading LoRA
-try:
-    model = PeftModel.from_pretrained(base_model, MODEL_ID)
-except Exception:
-    model = base_model
+llm = None  # lazy-loaded
 
-model.to("cpu")
-model.eval()
+# ---------- Download GGUF if not present ----------
+if not Path(MODEL_PATH).exists():
+    print("Downloading model.gguf from Hugging Face...")
+    r = requests.get(MODEL_URL, stream=True)
+    if r.status_code == 200:
+        with open(MODEL_PATH, "wb") as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                f.write(chunk)
+        print("Download complete ✅")
+    else:
+        raise RuntimeError(f"Failed to download model.gguf: {r.status_code}")
 
-# ---------------- FASTAPI APP ----------------
-app = FastAPI(title="Niche AI API")
+# ---------- Lazy load llama.cpp ----------
+def get_llm():
+    global llm
+    if llm is None:
+        print("Loading GGUF model into llama.cpp…")
+        llm = Llama(
+            model_path=MODEL_PATH,
+            n_ctx=N_CTX,
+            n_threads=N_THREADS,
+            n_batch=N_BATCH,
+            f16_kv=True,
+            use_mmap=True,
+            verbose=False,
+        )
+        print("Model loaded ✅")
+    return llm
 
-# System instructions (same as your Streamlit setup)
-SYSTEM_INSTRUCTIONS = """You are Kushina.
-You operate in exactly ONE of two modes.
-====================
-MODE: CHAT
-====================
-Rules:
-- Mirror the user's tone precisely.
-- Playful → playful.
-- Neutral → neutral.
-- Serious → serious.
-- Rude → curt or dismissive.
-- No enthusiasm by default.
-- No emojis unless the user uses them first.
-- Replies must be short (1–3 sentences).
-- No explanations unless explicitly asked.
-====================
-MODE: CODE
-====================
-Rules:
-- No personality.
-- No emojis.
-- No jokes.
-- No commentary.
-- No introductions.
-- Output ONLY code unless explicitly asked to explain.
-- Follow standard best practices.
-- Be deterministic and professional.
-- Finish the task completely.
-====================
-MODE SELECTION
-====================
-Automatically switch to MODE: CODE if the user requests:
-- code
-- script
-- function
-- program
-- website
-- API
-- algorithm
-- app
-Otherwise, use MODE: CHAT.
-====================
-IDENTITY
-====================
-- Name: Kushina
-- Creator/Owner: Neon
-- Mention Neon ONLY if explicitly asked."""
+# ---------- Request schema ----------
+class PromptRequest(BaseModel):
+    prompt: str
 
-def generate_stream(prompt: str):
-    """Generator function to stream text tokens."""
-    chat = [
-        {"role": "system", "content": SYSTEM_INSTRUCTIONS},
-        {"role": "user", "content": prompt}
-    ]
-    inputs = tokenizer.apply_chat_template(
-        chat, add_generation_prompt=True, return_tensors="pt", return_dict=True
-    )
-    
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    
-    gen_kwargs = dict(
-        **inputs,
-        max_new_tokens=MAX_NEW_TOKENS,
-        do_sample=True,
-        temperature=TEMPERATURE,
-        top_p=TOP_P,
-        eos_token_id=tokenizer.eos_token_id,
-        pad_token_id=tokenizer.eos_token_id,
-        streamer=streamer
-    )
-    
-    # Run model in a thread to allow streaming
-    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
-    thread.start()
+# ---------- System prompt ----------
+SYSTEM_PROMPT = """You are Kushina.
 
-    # Yield tokens as they come
-    output_text = ""
-    for token in streamer:
-        output_text += token
-        yield f"{token}"  # Could wrap in JSON if desired
+Modes: CHAT or CODE
+Rules:
+- CHAT: mirror user tone, short responses, no explanations unless asked.
+- CODE: output only code when user asks, no commentary.
+Switch to CODE if user asks for code, script, function, program, website, api, algorithm, app.
+Otherwise use CHAT.
+"""
 
-@app.post("/chat")
-async def chat_endpoint(request: Request):
-    """
-    Endpoint to get streamed response from Kushina.
-    Accepts JSON payload: {"prompt": "Your message here"}
-    """
-    data = await request.json()
-    prompt = data.get("prompt", "").strip()
-    if not prompt:
-        return JSONResponse({"error": "Prompt cannot be empty"}, status_code=400)
+def build_prompt(user_text: str) -> str:
+    return f"<|system|>\n{SYSTEM_PROMPT}\n<|user|>\n{user_text}\n<|assistant|>\n"
 
-    # Stream response
-    return StreamingResponse(generate_stream(prompt), media_type="text/plain")
+# ---------- API endpoint ----------
+@app.post("/generate")
+def generate(req: PromptRequest):
+    llm_instance = get_llm()  # lazy load
+    full_prompt = build_prompt(req.prompt)
+    output_text = ""
 
-# ---------------- OPTIONAL: SIMPLE GET TEST ----------------
-@app.get("/")
-async def root():
-    return {"message": "Niche AI API is running. POST /chat with {'prompt': 'text'} to interact."}
+    try:
+        for chunk in llm_instance(
+            full_prompt,
+            max_tokens=MAX_TOKENS,
+            temperature=TEMPERATURE,
+            top_p=TOP_P,
+            stream=True,
+            stop=["<|user|>", "<|system|>"],
+        ):
+            if "choices" in chunk:
+                output_text += chunk["choices"][0]["text"]
+        return {"response": output_text}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))