import spaces
import torch
from fastapi import HTTPException
from gradio import Server
from loguru import logger
from pydantic import BaseModel
from typing import List, Optional
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = "TrevorJS/gemma-4-E4B-it-uncensored"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = None

class ChatMessage(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = 512

@spaces.GPU
def run_inference(prompt: str, max_tokens: int = 512, temperature: float = 0.7) -> str:
    global model
    if model is None:
        logger.info(f"Loading model {MODEL_NAME}...")
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME, device_map="auto", torch_dtype=torch.bfloat16
        )
        logger.info("Model loaded")
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=temperature > 0
        )
    generated = outputs[0][inputs.input_ids.shape[-1]:]
    return tokenizer.decode(generated, skip_special_tokens=True)

app = Server()

@app.get("/")
def root():
    return {"message": "Gemma 4 E4B Uncensored API is running"}

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
    try:
        messages = [{"role": m.role, "content": m.content} for m in request.messages]
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        response = run_inference(prompt, request.max_tokens, request.temperature)
        return {
            "id": "chatcmpl-zerogpu",
            "object": "chat.completion",
            "model": request.model,
            "choices": [{
                "index": 0,
                "message": {"role": "assistant", "content": response},
                "finish_reason": "stop"
            }]
        }
    except Exception as e:
        logger.error(f"Error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

app.launch()