import gradio as gr
from transformers import pipeline
import torch
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
import uvicorn
import time

# 1. Load the model
model_id = "Aaresh5308a/Gemma4_e2b_NepalaiEdition"
pipe = pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device_map="auto",
)

# 2. Define the Inference Function
def chat_response(message, history):
    # Format the prompt for Gemma 4
    prompt = pipe.tokenizer.apply_chat_template(
        history + [{"role": "user", "content": message}],
        tokenize=False,
        add_generation_prompt=True
    )
    outputs = pipe(prompt, max_new_tokens=512, do_sample=True, temperature=0.7)
    return outputs[0]["generated_text"][len(prompt):]

# 3. Create Gradio UI
demo = gr.ChatInterface(fn=chat_response, title="Gemma 4 Nepali Edition")

# 4. Create OpenAI-Compatible API using FastAPI
app = FastAPI()

@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    data = await request.json()
    messages = data.get("messages", [])
    
    # Simple extraction of the last user message for this example
    user_msg = messages[-1]["content"]
    history = messages[:-1]
    
    response_text = chat_response(user_msg, history)
    
    # Return OpenAI formatted JSON
    return {
        "id": f"chatcmpl-{int(time.time())}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model_id,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": response_text},
            "finish_reason": "stop"
        }]
    }

# Mount Gradio to FastAPI
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)