Spaces:
Running on Zero
Running on Zero
| import os | |
| import torch | |
| from fastapi import FastAPI, HTTPException | |
| from pydantic import BaseModel | |
| from typing import List, Optional | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig | |
| import spaces | |
| import gradio as gr | |
| # Official Public Nous Research Repository (Kabhi delete ya 401 nahi hoga) | |
| MODEL_ID = "NousResearch/Hermes-3-Llama-3.1-8B" | |
| print("1. Loading Tokenizer...") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) | |
| model = None | |
| # 4-bit Quantization configuration for ZeroGPU (~5.5GB VRAM usage) | |
| bnb_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_use_double_quant=True, | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_compute_dtype=torch.float16 | |
| ) | |
| def get_or_load_model(): | |
| global model | |
| if model is None: | |
| print("2. Mounting 4-bit Model inside ZeroGPU context...") | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_ID, | |
| quantization_config=bnb_config, | |
| device_map="auto", | |
| low_cpu_mem_usage=True | |
| ) | |
| return model | |
| # Pydantic Schemas for OpenAI compatibility | |
| class Message(BaseModel): | |
| role: str | |
| content: str | |
| class ChatCompletionRequest(BaseModel): | |
| model: Optional[str] = "Hermes-3-8B" | |
| messages: List[Message] | |
| max_tokens: Optional[int] = 2048 | |
| temperature: Optional[float] = 0.6 | |
| # ZeroGPU Generation Function | |
| def generate_hermes(prompt_text: str, max_tokens: int, temperature: float) -> str: | |
| active_model = get_or_load_model() | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| inputs = tokenizer(prompt_text, return_tensors="pt").to(device) | |
| with torch.no_grad(): | |
| outputs = active_model.generate( | |
| **inputs, | |
| max_new_tokens=max_tokens, | |
| temperature=temperature, | |
| do_sample=True if temperature > 0.0 else False, | |
| top_p=0.9, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| generated_tokens = outputs[0][inputs.input_ids.shape[-1]:] | |
| # skip_special_tokens=False zaroori hai taaki <tool_call> tags strip na hon | |
| return tokenizer.decode(generated_tokens, skip_special_tokens=False) | |
| # UI element space ko active rakhne ke liye | |
| with gr.Blocks(title="Hermes-3 8B Server") as demo: | |
| gr.Markdown("# 🚀 Hermes-3 8B ZeroGPU API Server") | |
| gr.Markdown("Active Endpoints:\n- `GET /health`\n- `POST /v1/chat/completions`") | |
| app = demo.app | |
| def health_check(): | |
| return {"status": "ok", "ready": True} | |
| def chat_completions(req: ChatCompletionRequest): | |
| try: | |
| formatted_messages = [{"role": m.role, "content": m.content} for m in req.messages] | |
| # Hermes ChatML formatting | |
| prompt_text = tokenizer.apply_chat_template( | |
| formatted_messages, | |
| tokenize=False, | |
| add_generation_prompt=True | |
| ) | |
| raw_output = generate_hermes( | |
| prompt_text=prompt_text, | |
| max_tokens=req.max_tokens, | |
| temperature=req.temperature | |
| ) | |
| # Hermes stop tags trim karein | |
| clean_output = raw_output.replace("<|im_end|>", "").strip() | |
| return { | |
| "id": "chatcmpl-hermes", | |
| "object": "chat.completion", | |
| "model": req.model, | |
| "choices": [ | |
| { | |
| "index": 0, | |
| "message": { | |
| "role": "assistant", | |
| "content": clean_output | |
| }, | |
| "finish_reason": "stop" | |
| } | |
| ] | |
| } | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| if __name__ == "__main__": | |
| demo.queue().launch(server_name="0.0.0.0", server_port=7860) | |