import argparse import asyncio import json import os import sys import threading import time from typing import List, Dict, Any, Union from fastapi import FastAPI, Request from fastapi.responses import StreamingResponse from huggingface_hub import HfApi from pydantic import BaseModel import torch import uvicorn from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer app = FastAPI(title="Gemma-4 HF API Server") # Global model & tokenizer references model = None tokenizer = None loaded_repo_id = None WEIGHT_FILENAMES = { "model.safetensors", "model.safetensors.index.json", "pytorch_model.bin", "pytorch_model.bin.index.json", } class ChatMessage(BaseModel): role: str content: str class ChatCompletionRequest(BaseModel): model: str = "OBLITERATUS/Gemma-4-12B-OBLITERATED" messages: List[ChatMessage] temperature: float = 0.7 top_p: float = 0.9 top_k: int = 40 max_tokens: int = 512 stream: bool = False repetition_penalty: float = 1.1 def repo_has_transformers_weights(repo_id: str, token: str | None) -> bool: files = HfApi(token=token).list_repo_files(repo_id=repo_id, repo_type="model") return any( filename in WEIGHT_FILENAMES or filename.endswith(".safetensors") or filename.endswith(".bin") for filename in files ) def resolve_repo_id( repo_id: str, fallback_repo_id: str | None, wait_for_weights: int, poll_interval: int, token: str | None, ) -> str: deadline = time.monotonic() + wait_for_weights while True: if repo_has_transformers_weights(repo_id, token): return repo_id if time.monotonic() >= deadline: break remaining = max(0, int(deadline - time.monotonic())) print( f"No Transformers weights are published for {repo_id} yet. " f"Checking again in {poll_interval}s ({remaining}s remaining)...", flush=True, ) time.sleep(min(poll_interval, remaining)) if fallback_repo_id: if not repo_has_transformers_weights(fallback_repo_id, token): raise RuntimeError( f"Neither {repo_id} nor fallback {fallback_repo_id} contains " "Transformers weights." ) print( f"WARNING: {repo_id} has no Transformers weights. " f"Using the explicitly requested fallback {fallback_repo_id}.", flush=True, ) return fallback_repo_id raise RuntimeError( f"{repo_id} does not currently contain model weights. Its Hugging Face " "repository only publishes configuration/tokenizer files, so " "AutoModelForCausalLM cannot load it.\n" "Wait for the advertised weight files to finish publishing, or run an " "explicit fallback, for example:\n" " --fallback-repo-id google/gemma-4-12B-it\n" "To wait for an in-progress upload, add:\n" " --wait-for-weights 3600" ) def load_model( repo_id: str, fallback_repo_id: str | None = None, wait_for_weights: int = 0, poll_interval: int = 60, ): global model, tokenizer, loaded_repo_id token = os.environ.get("HF_TOKEN") or None selected_repo_id = resolve_repo_id( repo_id=repo_id, fallback_repo_id=fallback_repo_id, wait_for_weights=max(0, wait_for_weights), poll_interval=max(5, poll_interval), token=token, ) if token is None: print( "HF_TOKEN is not set. Public downloads still work, but Hugging Face " "applies lower rate limits.", flush=True, ) print(f"Loading tokenizer for {selected_repo_id}...") tokenizer = AutoTokenizer.from_pretrained( selected_repo_id, trust_remote_code=True, token=token, ) print(f"Loading model weights for {selected_repo_id} (device_map='auto')...") model = AutoModelForCausalLM.from_pretrained( selected_repo_id, device_map="auto", torch_dtype="auto", trust_remote_code=True, token=token, ) loaded_repo_id = selected_repo_id print(f"Model loaded successfully: {selected_repo_id}") @app.get("/health") async def health(): return { "status": "ok" if model is not None and tokenizer is not None else "loading", "model": loaded_repo_id, } async def stream_generator(streamer: TextIteratorStreamer): loop = asyncio.get_event_loop() while True: try: token = await loop.run_in_executor(None, lambda: next(streamer, None)) if token is None: break chunk = { "choices": [ { "delta": {"content": token}, "finish_reason": None, "index": 0 } ] } yield f"data: {json.dumps(chunk)}\n\n" except Exception as e: print(f"Error in stream: {e}") break chunk_done = { "choices": [ { "delta": {}, "finish_reason": "stop", "index": 0 } ] } yield f"data: {json.dumps(chunk_done)}\n\n" yield "data: [DONE]\n\n" @app.post("/v1/chat/completions") async def chat_completions(request: ChatCompletionRequest): global model, tokenizer if model is None or tokenizer is None: return {"error": "Model not loaded"} messages_list = [{"role": msg.role, "content": msg.content} for msg in request.messages] prompt = tokenizer.apply_chat_template( messages_list, tokenize=False, add_generation_prompt=True, enable_thinking=False, ) inputs = tokenizer(prompt, return_tensors="pt").to(model.device) if request.stream: streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) generation_kwargs = dict( **inputs, max_new_tokens=request.max_tokens, temperature=request.temperature, top_p=request.top_p, top_k=request.top_k, do_sample=True, repetition_penalty=request.repetition_penalty, streamer=streamer, ) thread = threading.Thread(target=model.generate, kwargs=generation_kwargs) thread.start() return StreamingResponse(stream_generator(streamer), media_type="text/event-stream") else: with torch.no_grad(): output = model.generate( **inputs, max_new_tokens=request.max_tokens, temperature=request.temperature, top_p=request.top_p, top_k=request.top_k, do_sample=True, repetition_penalty=request.repetition_penalty, ) generated_text = tokenizer.decode(output[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True) return { "choices": [ { "message": { "role": "assistant", "content": generated_text }, "finish_reason": "stop", "index": 0 } ] } if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--repo-id", type=str, default="OBLITERATUS/Gemma-4-12B-OBLITERATED") parser.add_argument( "--fallback-repo-id", type=str, default=None, help="Explicit model to use only when --repo-id has no published weights.", ) parser.add_argument( "--wait-for-weights", type=int, default=0, metavar="SECONDS", help="Wait for an in-progress Hugging Face upload before failing.", ) parser.add_argument( "--poll-interval", type=int, default=60, metavar="SECONDS", help="Hugging Face polling interval used with --wait-for-weights.", ) parser.add_argument("--host", type=str, default="0.0.0.0") parser.add_argument("--port", type=int, default=8000) args = parser.parse_args() try: load_model( repo_id=args.repo_id, fallback_repo_id=args.fallback_repo_id, wait_for_weights=args.wait_for_weights, poll_interval=args.poll_interval, ) except RuntimeError as exc: print(f"\nERROR: {exc}", file=sys.stderr) raise SystemExit(2) from None uvicorn.run(app, host=args.host, port=args.port)