Voxtreme-ai / app.py
ppgodzilla's picture
Update app.py
c23131f verified
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})