scratch/serve_gemma.py

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)