Update app.py

app.py

@@ -1,27 +1,43 @@
 import os
+import sys
 import secrets
 import spaces
 import torch
 import gradio as gr
-from typing import List, Union, Optional
+from typing import List, Union, Optional, Dict, Any
 from fastapi import Header, HTTPException, status
 from pydantic import BaseModel
+from huggingface_hub import snapshot_download
 from sentence_transformers import SentenceTransformer
 
-MODEL_ID = "Dibachain/Diba-Embed"
+EMBED_MODEL_ID = "Dibachain/Diba-Embed"
+MIZAN_MODEL_ID = "Nimruz/Mizan-v3-3B"
 API_KEY = os.environ.get("MY_API_KEY", "")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-print(f"Loading {MODEL_ID}...")
-model = SentenceTransformer(
-    MODEL_ID,
+# 1. Load Diba-Embed
+print(f"Loading {EMBED_MODEL_ID} on {device}...")
+embed_model = SentenceTransformer(
+    EMBED_MODEL_ID,
     trust_remote_code=True,
     device=device,
     model_kwargs={"torch_dtype": torch.bfloat16}
 )
-model.max_seq_length = 512
-print("Model loaded successfully!")
+embed_model.max_seq_length = 512
+print("Diba-Embed loaded successfully!")
 
+# 2. Download and Load Nimruz/Mizan-v3-3B via its native joint_schema_model.py
+print(f"Downloading and loading {MIZAN_MODEL_ID} on {device}...")
+mizan_path = snapshot_download(MIZAN_MODEL_ID)
+if mizan_path not in sys.path:
+    sys.path.insert(0, mizan_path)
+
+from joint_schema_model import load_release_model, systemone
+
+mizan_model, mizan_processor = load_release_model(mizan_path, device=device)
+print("Mizan-v3-3B loaded successfully!")
+
+# 3. API Key Helper
 def verify_key(provided_key: Optional[str]):
     if not API_KEY:
         raise RuntimeError("MY_API_KEY secret is not configured in Space Settings.")
@@ -31,7 +47,7 @@ def verify_key(provided_key: Optional[str]):
             detail="Invalid or missing API Key"
         )
 
-# 1. ZeroGPU function protected with an internal key check so random visitors can't use the UI
+# 4. ZeroGPU Function 1: Embeddings
 @spaces.GPU(duration=30)
 def compute_embeddings(
     api_key: str,
@@ -45,7 +61,7 @@ def compute_embeddings(
     text_list = [texts] if isinstance(texts, str) else texts
     prompt_name = "query" if input_type == "query" else None
 
-    embeddings = model.encode(
+    embeddings = embed_model.encode(
         text_list,
         prompt_name=prompt_name,
         normalize_embeddings=normalize,
@@ -56,47 +72,109 @@ def compute_embeddings(
         for i, emb in enumerate(embeddings)
     ]
 
-# 2. Gradio Interface (api_name=False hides it from public Gradio API clients)
-demo = gr.Interface(
+# 5. ZeroGPU Function 2: Mizan-v3-3B Decision Inference
+@spaces.GPU(duration=30)
+def compute_decision(
+    api_key: str,
+    state: Any,
+    questions: Dict[str, Any]
+):
+    if not API_KEY or not secrets.compare_digest(api_key or "", API_KEY):
+        raise gr.Error("Unauthorized: Invalid API Key")
+
+    payload = {
+        "model": MIZAN_MODEL_ID,
+        "state": state,
+        "questions": questions
+    }
+    with torch.no_grad():
+        return systemone(mizan_model, mizan_processor, payload)
+
+# 6. Build a Tabbed Gradio Interface so ZeroGPU scans and registers BOTH @spaces.GPU functions
+embed_ui = gr.Interface(
     fn=compute_embeddings,
     inputs=[
-        gr.Textbox(label="API Key", type="password", placeholder="Enter MY_API_KEY..."),
-        gr.Textbox(label="Input Text", placeholder="Enter text to embed..."),
+        gr.Textbox(label="API Key", type="password"),
+        gr.Textbox(label="Input Text"),
         gr.Radio(["document", "query"], value="document", label="Input Type"),
         gr.Checkbox(value=True, label="Normalize")
     ],
-    outputs=gr.JSON(label="Output"),
-    title="Diba-Embed Private API",
+    outputs=gr.JSON(label="Embeddings Output"),
+    api_name=False
+)
+
+mizan_ui = gr.Interface(
+    fn=compute_decision,
+    inputs=[
+        gr.Textbox(label="API Key", type="password"),
+        gr.Textbox(label="State (Text or JSON)"),
+        gr.JSON(label="Questions Dict")
+    ],
+    outputs=gr.JSON(label="Decision Output"),
     api_name=False
 )
 
-# 3. Launch Gradio with show_api=False to disable public /gradio_api docs
+demo = gr.TabbedInterface(
+    [embed_ui, mizan_ui],
+    ["Diba-Embed", "Mizan-v3-3B"],
+    title="Diba-Embed & Mizan-v3-3B Private API"
+)
+
+# Launch Gradio non-blocking on port 7860
 demo.launch(prevent_thread_lock=True, show_api=False)
 
-# 4. Protected FastAPI /v1/embeddings endpoint
+# 7. FastAPI Schemas & Custom Endpoints attached to demo.app
 class EmbedRequest(BaseModel):
     input: Union[str, List[str]]
     input_type: Optional[str] = "document"
     normalize: Optional[bool] = True
 
+class DecisionRequest(BaseModel):
+    model: Optional[str] = MIZAN_MODEL_ID
+    state: Any
+    questions: Dict[str, Any]
+
+def extract_token(authorization: Optional[str], x_api_key: Optional[str]) -> Optional[str]:
+    if x_api_key:
+        return x_api_key
+    if authorization and authorization.lower().startswith("bearer "):
+        return authorization[7:].strip()
+    return None
+
+@demo.app.get("/health")
+def health_check():
+    return {
+        "status": "online",
+        "models": [EMBED_MODEL_ID, MIZAN_MODEL_ID],
+        "endpoints": ["/v1/embeddings", "/v1/systemone", "/v1/decide"]
+    }
+
 @demo.app.post("/v1/embeddings")
 def create_embeddings(
     req: EmbedRequest,
     authorization: Optional[str] = Header(default=None),
     x_api_key: Optional[str] = Header(default=None)
 ):
-    # Accept either "Authorization: Bearer <MY_API_KEY>" or "X-API-Key: <MY_API_KEY>"
-    token = x_api_key
-    if not token and authorization and authorization.lower().startswith("bearer "):
-        token = authorization[7:].strip()
-
+    token = extract_token(authorization, x_api_key)
     verify_key(token)
 
     data = compute_embeddings(token, req.input, req.input_type or "document", req.normalize)
     return {
         "object": "list",
-        "model": MODEL_ID,
+        "model": EMBED_MODEL_ID,
         "data": data
     }
 
+@demo.app.post("/v1/systemone")
+@demo.app.post("/v1/decide")
+def create_decision(
+    req: DecisionRequest,
+    authorization: Optional[str] = Header(default=None),
+    x_api_key: Optional[str] = Header(default=None)
+):
+    token = extract_token(authorization, x_api_key)
+    verify_key(token)
+
+    return compute_decision(token, req.state, req.questions)
+
 demo.block_thread()