Spaces:
Running on Zero
Running on Zero
| import os | |
| import sys | |
| import secrets | |
| import spaces | |
| import torch | |
| import gradio as gr | |
| from typing import List, Union, Optional, Dict, Any | |
| from fastapi import APIRouter, Header, HTTPException, status | |
| from pydantic import BaseModel | |
| from huggingface_hub import snapshot_download | |
| from sentence_transformers import SentenceTransformer | |
| 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" | |
| # 1. Load Diba-Embed and register with ZeroGPU | |
| print(f"Loading {EMBED_MODEL_ID}...") | |
| embed_model = SentenceTransformer( | |
| EMBED_MODEL_ID, | |
| trust_remote_code=True, | |
| model_kwargs={"torch_dtype": torch.bfloat16} | |
| ) | |
| embed_model.max_seq_length = 512 | |
| if torch.cuda.is_available(): | |
| embed_model = embed_model.to("cuda") | |
| print("Diba-Embed loaded successfully!") | |
| # 2. Download and Load Nimruz/Mizan-v3-3B | |
| print(f"Downloading and loading {MIZAN_MODEL_ID}...") | |
| 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 | |
| target_device = "cuda" if torch.cuda.is_available() else "cpu" | |
| mizan_model, mizan_processor = load_release_model(mizan_path, device=target_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.") | |
| if not provided_key or not secrets.compare_digest(provided_key, API_KEY): | |
| raise HTTPException( | |
| status_code=status.HTTP_401_UNAUTHORIZED, | |
| detail="Invalid or missing API Key" | |
| ) | |
| # 4. ZeroGPU Function 1: Embeddings | |
| def compute_embeddings( | |
| api_key: str, | |
| texts: Union[str, List[str]], | |
| input_type: str = "document", | |
| normalize: bool = True | |
| ): | |
| if not API_KEY or not secrets.compare_digest(api_key or "", API_KEY): | |
| raise gr.Error("Unauthorized: Invalid API Key") | |
| text_list = [texts] if isinstance(texts, str) else texts | |
| prompt_name = "query" if input_type == "query" else None | |
| embeddings = embed_model.encode( | |
| text_list, | |
| prompt_name=prompt_name, | |
| normalize_embeddings=normalize, | |
| batch_size=8 | |
| ) | |
| return [ | |
| {"object": "embedding", "index": i, "embedding": emb.tolist()} | |
| for i, emb in enumerate(embeddings) | |
| ] | |
| # 5. ZeroGPU Function 2: Mizan-v3-3B Decision Inference | |
| 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"), | |
| gr.Textbox(label="Input Text"), | |
| gr.Radio(["document", "query"], value="document", label="Input Type"), | |
| gr.Checkbox(value=True, label="Normalize") | |
| ], | |
| 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 | |
| ) | |
| demo = gr.TabbedInterface( | |
| [embed_ui, mizan_ui], | |
| ["Diba-Embed", "Mizan-v3-3B"], | |
| title="Diba-Embed & Mizan-v3-3B Private API" | |
| ) | |
| # Launch Gradio with ssr_mode=False so FastAPI binds directly to port 7860 | |
| demo.launch(prevent_thread_lock=True, ssr_mode=False) | |
| # 7. Create an APIRouter for custom endpoints | |
| api_router = APIRouter() | |
| 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 | |
| def health_check(): | |
| return { | |
| "status": "online", | |
| "models": [EMBED_MODEL_ID, MIZAN_MODEL_ID], | |
| "endpoints": ["/v1/embeddings", "/v1/systemone", "/v1/decide"] | |
| } | |
| def create_embeddings( | |
| req: EmbedRequest, | |
| authorization: Optional[str] = Header(default=None), | |
| x_api_key: Optional[str] = Header(default=None) | |
| ): | |
| 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": EMBED_MODEL_ID, | |
| "data": data | |
| } | |
| 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) | |
| # Prepend custom routes BEFORE Gradio's catch-all GET /{path:path} route | |
| demo.app.router.routes[:0] = api_router.routes | |
| demo.block_thread() |