Spaces:
Running
Running
File size: 1,302 Bytes
253c711 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | 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}
|