Download server.py from Flux91/Swastik: direct link, hf CLI and curl.
- Browser
- Download file 17.9 kB
-
https://huggingface.co/spaces/Flux91/Swastik/resolve/main/server.py
- Command line
-
hf download hf://spaces/Flux91/Swastik/server.py
-
curl -L -o server.py https://huggingface.co/spaces/Flux91/Swastik/resolve/main/server.py
17.9 kB
| #!/usr/bin/env python3 | |
| """ | |
| server.py β High-performance FastAPI backend & Modern Web UI for Mac Local Inference. | |
| Provides: | |
| - REST API endpoint (/api/chat) | |
| - Real-time token streaming endpoint (/api/chat/stream) via Server-Sent Events | |
| - System status endpoint (/api/status) | |
| - Built-in Swagger docs (/docs) | |
| - Serves modern, glassmorphic Web UI (/) | |
| """ | |
| import os | |
| import sys | |
| import json | |
| import time | |
| import base64 | |
| import asyncio | |
| import threading | |
| from typing import List, Optional, Dict, Any | |
| from pathlib import Path | |
| from fastapi import FastAPI, Request, HTTPException | |
| from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse, FileResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel | |
| # Import MLX loader and generator from chat module | |
| from chat import SYSTEM_PROMPT, try_mlx, try_transformers | |
| # Voice-to-Voice engine (Whisper STT + Piper TTS) | |
| from voice import engine as voice_engine | |
| # Instruct the brain to answer in a spoken, TTS-friendly way during voice mode | |
| VOICE_INSTRUCTION = ( | |
| "You are speaking to the user aloud by voice right now. " | |
| "Answer as a natural spoken conversation: be concise, friendly, and clear. " | |
| "Use short, simple sentences. Plain text only β NO markdown, code blocks, " | |
| "headings, bullet points, emojis, or special symbols. " | |
| "IMPORTANT: Reply in the SAME language the user just spoke in " | |
| "(e.g. if they speak Hindi, answer in Hindi)." | |
| ) | |
| # Global model state | |
| backend = None | |
| model_info = { | |
| "name": "Loading...", | |
| "backend": "none", | |
| "path": "", | |
| "status": "initializing" | |
| } | |
| BASE_DIR = Path(__file__).resolve().parent | |
| STATIC_DIR = BASE_DIR / "static" | |
| MODELS_DIR = BASE_DIR / "models" / "swastik-0.5b" | |
| def load_active_model(): | |
| """Load the best available local model.""" | |
| global backend, model_info | |
| target_path = None | |
| model_name = "Flux91/swastik-0.5b" | |
| if MODELS_DIR.is_dir() and (MODELS_DIR / "model.safetensors").exists(): | |
| target_path = str(MODELS_DIR) | |
| model_name = "Flux91/swastik-0.5b (Fine-Tuned MLX 4-bit)" | |
| else: | |
| # Fallback to remote or default | |
| target_path = "Flux91/swastik-0.5b" | |
| model_name = "Flux91/swastik-0.5b" | |
| print(f"π Initializing model from: {target_path}") | |
| b, _ = try_mlx(target_path) | |
| if b is None: | |
| print("β οΈ MLX failed, attempting transformers fallback...") | |
| b, _ = try_transformers(target_path) | |
| if b is not None: | |
| backend = b | |
| model_info = { | |
| "name": model_name, | |
| "backend": b[0].upper(), | |
| "path": target_path, | |
| "status": "ready" | |
| } | |
| print(f"β Model loaded successfully via {model_info['backend']}!") | |
| else: | |
| model_info["status"] = "error" | |
| print("β Could not load model.") | |
| # Initialize model at startup | |
| load_active_model() | |
| app = FastAPI( | |
| title="AI Chatbot API", | |
| description="High-performance local inference powered by Apple Silicon MLX", | |
| version="1.0.0" | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # Mount static folder | |
| STATIC_DIR.mkdir(parents=True, exist_ok=True) | |
| app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static") | |
| class MessageItem(BaseModel): | |
| role: str | |
| content: Any | |
| class ChatRequest(BaseModel): | |
| message: str | |
| history: Optional[List[MessageItem]] = [] | |
| temperature: Optional[float] = 0.7 | |
| max_tokens: Optional[int] = 512 | |
| system_prompt: Optional[str] = None | |
| def build_sanitized_messages(message: str, history: List[MessageItem], system_prompt: Optional[str] = None): | |
| """Build a sanitized list of {role, content: str} messages. No system prompt injected by default.""" | |
| messages = [] | |
| # Only add system prompt if user explicitly provided one in settings | |
| if system_prompt and system_prompt.strip(): | |
| messages.append({"role": "system", "content": system_prompt.strip()}) | |
| if history: | |
| for item in history: | |
| role = item.role | |
| c = item.content | |
| if isinstance(c, list): | |
| c = "".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in c) | |
| elif isinstance(c, dict): | |
| c = c.get("text", str(c)) | |
| messages.append({"role": role, "content": str(c)}) | |
| user_c = message | |
| if isinstance(user_c, list): | |
| user_c = "".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in user_c) | |
| elif isinstance(user_c, dict): | |
| user_c = user_c.get("text", str(user_c)) | |
| messages.append({"role": "user", "content": str(user_c)}) | |
| return messages | |
| async def index(): | |
| """Serve the main single page web application.""" | |
| index_file = STATIC_DIR / "index.html" | |
| if index_file.exists(): | |
| return FileResponse(str(index_file)) | |
| return HTMLResponse("<h1>AI Chatbot API is running. Missing index.html in static/</h1>") | |
| async def get_status(): | |
| """Return current server and model status.""" | |
| return { | |
| "status": model_info["status"], | |
| "model_name": model_info["name"], | |
| "backend": model_info["backend"], | |
| "path": model_info["path"], | |
| "apple_silicon": True | |
| } | |
| async def chat_endpoint(req: ChatRequest): | |
| """Standard JSON API chat endpoint.""" | |
| global backend | |
| if backend is None: | |
| raise HTTPException(status_code=503, detail="Model is not loaded.") | |
| backend_type = backend[0] | |
| messages = build_sanitized_messages(req.message, req.history, req.system_prompt) | |
| start_t = time.time() | |
| if backend_type == "mlx": | |
| from chat import generate_mlx | |
| reply = generate_mlx(backend, messages, req.max_tokens, req.temperature) | |
| else: | |
| from chat import generate_transformers | |
| reply = generate_transformers(backend, messages, req.max_tokens, req.temperature) | |
| duration = time.time() - start_t | |
| return { | |
| "reply": reply, | |
| "model": model_info["name"], | |
| "backend": model_info["backend"], | |
| "latency_sec": round(duration, 3) | |
| } | |
| async def chat_stream_endpoint(req: ChatRequest): | |
| """Real-time token streaming endpoint using Server-Sent Events (SSE).""" | |
| global backend | |
| if backend is None: | |
| raise HTTPException(status_code=503, detail="Model is not loaded.") | |
| backend_type = backend[0] | |
| messages = build_sanitized_messages(req.message, req.history, req.system_prompt) | |
| if backend_type != "mlx": | |
| # Real-time token streaming for PyTorch / Transformers | |
| from chat import stream_transformers | |
| async def trans_stream(): | |
| loop = asyncio.get_event_loop() | |
| token_queue = asyncio.Queue() | |
| def run_generation(): | |
| try: | |
| for tok in stream_transformers(backend, messages, req.max_tokens, req.temperature): | |
| loop.call_soon_threadsafe(token_queue.put_nowait, tok) | |
| except Exception as e: | |
| print(f"β οΈ Streaming error in transformers backend: {e}") | |
| finally: | |
| loop.call_soon_threadsafe(token_queue.put_nowait, None) | |
| threading.Thread(target=run_generation, daemon=True).start() | |
| while True: | |
| tok = await token_queue.get() | |
| if tok is None: | |
| break | |
| payload = json.dumps({"token": tok}) | |
| yield f"data: {payload}\n\n" | |
| yield "data: [DONE]\n\n" | |
| return StreamingResponse( | |
| trans_stream(), | |
| media_type="text/event-stream", | |
| headers={ | |
| "Cache-Control": "no-cache", | |
| "Connection": "keep-alive", | |
| "X-Accel-Buffering": "no" | |
| } | |
| ) | |
| # Real MLX token streaming | |
| _, model, tokenizer = backend | |
| if hasattr(tokenizer, 'apply_chat_template'): | |
| prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
| else: | |
| prompt = "".join(f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n" for m in messages) + "<|im_start|>assistant\n" | |
| from mlx_lm.generate import stream_generate | |
| from mlx_lm.sample_utils import make_sampler | |
| sampler = make_sampler(temp=req.temperature) | |
| async def token_generator(): | |
| # Run generator in separate thread to prevent blocking the async event loop | |
| loop = asyncio.get_event_loop() | |
| def generate_all_tokens(): | |
| tokens = [] | |
| for chunk in stream_generate(model, tokenizer, prompt, max_tokens=req.max_tokens, sampler=sampler): | |
| tokens.append(chunk.text) | |
| return tokens | |
| # Yield tokens as generated | |
| for chunk in stream_generate(model, tokenizer, prompt, max_tokens=req.max_tokens, sampler=sampler): | |
| payload = json.dumps({"token": chunk.text}) | |
| yield f"data: {payload}\n\n" | |
| await asyncio.sleep(0.005) # Yield control to event loop | |
| yield "data: [DONE]\n\n" | |
| return StreamingResponse( | |
| token_generator(), | |
| media_type="text/event-stream", | |
| headers={ | |
| "Cache-Control": "no-cache", | |
| "Connection": "keep-alive", | |
| "X-Accel-Buffering": "no" | |
| } | |
| ) | |
| # ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| # VOICE-TO-VOICE ENDPOINTS | |
| # ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| async def voice_status(): | |
| """Voice components status + download progress.""" | |
| s = voice_engine.status() | |
| s["piper_available"] = voice_engine._piper is not None | |
| s["say_available"] = voice_engine._say_ok | |
| return s | |
| async def voice_init(): | |
| """Start downloading/loading voice models (idempotent, returns immediately).""" | |
| voice_engine.start_init() | |
| return voice_engine.status() | |
| class VoiceRoundRequest(BaseModel): | |
| audio: str # base64 WAV (16-bit PCM, 16 kHz mono) | |
| history: Optional[List[MessageItem]] = [] | |
| system_prompt: Optional[str] = None | |
| temperature: Optional[float] = 0.7 | |
| max_tokens: Optional[int] = 256 | |
| async def voice_round(req: VoiceRoundRequest, request: Request): | |
| """One voice round: STT -> LLM -> TTS, streamed as Server-Sent Events. | |
| Each `data:` line is a JSON object with a `type` field: | |
| loading β "loading" events during model download/loading | |
| sts β {text, language} Whisper English transcript | |
| ai_start β LLM generation started | |
| token β {text} live token for the overlay transcript | |
| audio β {data} base64 WAV (sentence-sized) spoken via TTS | |
| ai_end β {text} full assistant reply | |
| done β stream complete | |
| error β {message} | |
| """ | |
| global backend | |
| try: | |
| wav_bytes = base64.b64decode(req.audio) | |
| except Exception: | |
| raise HTTPException(status_code=422, detail="Invalid base64 audio.") | |
| if backend is None: | |
| raise HTTPException(status_code=503, detail="LLM not loaded.") | |
| LANG_NAMES = { | |
| "en": "English", "hi": "Hindi", "ta": "Tamil", "bn": "Bengali", | |
| "mr": "Marathi", "te": "Telugu", "kn": "Kannada", "gu": "Gujarati", | |
| "ur": "Urdu", "pa": "Punjabi", "ml": "Malayalam", "or": "Odia", | |
| "as": "Assamese", "ne": "Nepali", "hi-en": "Hindi", | |
| } | |
| def build_voice_messages(user_text, user_history, lang_code="en"): | |
| lang_name = LANG_NAMES.get((lang_code or "en").lower().split("-")[0]) | |
| msgs = [] | |
| sys_parts = [] | |
| if req.system_prompt and req.system_prompt.strip(): | |
| sys_parts.append(req.system_prompt.strip()) | |
| sys_parts.append(VOICE_INSTRUCTION) | |
| if lang_name and lang_name != "English": | |
| sys_parts.append( | |
| f"The user just spoke to you in {lang_name} (language code '{lang_code}'). " | |
| f"Your ENTIRE spoken reply must be in {lang_name} β same language and script " | |
| f"as the user. Do NOT switch to English. Answer naturally, in a sentence or two." | |
| ) | |
| msgs.append({"role": "system", "content": "\n\n".join(sys_parts)}) | |
| for item in user_history or []: | |
| c = item.content | |
| if isinstance(c, list): | |
| c = "".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in c) | |
| elif isinstance(c, dict): | |
| c = c.get("text", str(c)) | |
| msgs.append({"role": item.role, "content": str(c)}) | |
| msgs.append({"role": "user", "content": user_text}) | |
| return msgs | |
| async def event_stream(): | |
| def sse(payload): | |
| return f"data: {json.dumps(payload)}\n\n" | |
| loop = asyncio.get_event_loop() | |
| queue = asyncio.Queue() | |
| def worker(): | |
| def put(obj): | |
| loop.call_soon_threadsafe(queue.put_nowait, obj) | |
| try: | |
| # 1. Ensure voice models loaded | |
| if not voice_engine.is_ready(): | |
| voice_engine.ensure_loaded(progress_cb=lambda st: put({"type": "loading", **st})) | |
| if not voice_engine.is_ready(): | |
| put({"type": "error", "message": "Voice models could not be loaded."}) | |
| return | |
| # 2. STT (Whisper β keeps the user's own language) | |
| try: | |
| stt = voice_engine.transcribe(wav_bytes) | |
| except Exception as e: | |
| put({"type": "error", "message": "Transcription failed: " + str(e)}) | |
| return | |
| user_text = (stt.get("text") or "").strip() | |
| language = stt.get("language") or "en" | |
| if not user_text: | |
| put({"type": "error", "message": "Could not understand the audio."}) | |
| return | |
| put({"type": "sts", "text": user_text, "language": language}) | |
| # Ensure the TTS voice for the detected language is available | |
| try: | |
| voice_engine.ensure_lang(language, lambda st: put({"type": "loading", **st})) | |
| except Exception as e: | |
| print(f" [voice] could not load voice for language '{language}': {e}") | |
| # 3. LLM generation + per-sentence TTS | |
| put({"type": "ai_start"}) | |
| messages = build_voice_messages(user_text, req.history or [], language) | |
| b_type = backend[0] | |
| if b_type == "mlx": | |
| from mlx_lm.sample_utils import make_sampler | |
| from mlx_lm.generate import stream_generate | |
| _, model, tokenizer = backend | |
| prompt = tokenizer.apply_chat_template( | |
| messages, tokenize=False, add_generation_prompt=True | |
| ) | |
| sampler = make_sampler(temp=req.temperature) | |
| full = "" | |
| sentence = "" | |
| for chunk in stream_generate( | |
| model, tokenizer, prompt, | |
| max_tokens=req.max_tokens, sampler=sampler, | |
| ): | |
| token = chunk.text | |
| full += token | |
| sentence += token | |
| put({"type": "token", "text": token}) | |
| if sentence.strip() and (is_sentence_end(sentence) or len(sentence) >= 180): | |
| for wav_b64 in voice_engine.synth_chunks(sentence, language): | |
| put({"type": "audio", "data": wav_b64, "language": language}) | |
| sentence = "" | |
| if sentence.strip(): | |
| for wav_b64 in voice_engine.synth_chunks(sentence, language): | |
| put({"type": "audio", "data": wav_b64, "language": language}) | |
| put({"type": "ai_end", "text": full}) | |
| put({"type": "done"}) | |
| else: | |
| put({"type": "error", "message": "Voice mode requires the MLX backend."}) | |
| except Exception as e: | |
| import traceback | |
| traceback.print_exc() | |
| put({"type": "error", "message": str(e)}) | |
| threading.Thread(target=worker, daemon=True).start() | |
| while True: | |
| if await request.is_disconnected(): | |
| break | |
| item = await queue.get() | |
| yield sse(item) | |
| if item.get("type") in ("done", "error"): | |
| break | |
| return StreamingResponse( | |
| event_stream(), | |
| media_type="text/event-stream", | |
| headers={ | |
| "Cache-Control": "no-cache", | |
| "Connection": "keep-alive", | |
| "X-Accel-Buffering": "no" | |
| } | |
| ) | |
| def is_sentence_end(text): | |
| """True if text ends with a sentence-terminating punctuation.""" | |
| return text.rstrip().endswith((".", "!", "?", "ΰ₯€", "γ", "οΌ", "οΌ")) | |
| if __name__ == "__main__": | |
| import uvicorn | |
| default_port = 7860 if os.getenv("SPACE_ID") or os.getenv("HF_SPACE") else 8000 | |
| port = int(os.getenv("PORT", default_port)) | |
| print(f"\nπ AI Chatbot Server running at http://localhost:{port}") | |
| print(f"π Interactive API Documentation at http://localhost:{port}/docs\n") | |
| uvicorn.run(app, host="0.0.0.0", port=port, log_level="info") | |