from fastapi import FastAPI, Request
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# ✅ Initialize FastAPI app
app = FastAPI(title="Milo AI Backend", version="1.0")

# ✅ Load model and tokenizer once at startup
model_name = "milo-phi3"
device = "cuda" if torch.cuda.is_available() else "cpu"

model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# ✅ Define request schema for clarity
class MiloRequest(BaseModel):
    message: str

@app.post("/milo")
async def milo_generate(request: MiloRequest):
    """
    Generate a WhatsApp-style reply using Milo AI.
    """
    # Tokenize input
    inputs = tokenizer(request.message, return_tensors="pt").to(device)

    # Generate response
    output = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=True,          # Enable sampling for natural replies
        temperature=0.7,         # Controls creativity
        top_p=0.9,               # Nucleus sampling for variety
        pad_token_id=tokenizer.eos_token_id
    )

    # Decode and clean response
    reply = tokenizer.decode(output[0], skip_special_tokens=True)

    return {"reply": reply}