from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from transformers import pipeline
import uvicorn

app = FastAPI()

# Load your model
generator = pipeline(
    "text-generation",
    model="mistralai/Mistral-7B-Instruct",
    max_new_tokens=400,
    temperature=0.7,
    do_sample=True
)

@app.post("/ask")
async def ask_voxtreme(request: Request):
    data = await request.json()
    prompt = data.get("prompt", "")

    if not prompt:
        return JSONResponse(content={"error": "Prompt is missing"}, status_code=400)

    # Generate text using model
    output = generator(prompt)[0]["generated_text"]
    return JSONResponse(content={"response": output})