|
from fastapi import FastAPI, Request |
|
from fastapi.responses import JSONResponse |
|
from transformers import pipeline |
|
import uvicorn |
|
|
|
app = FastAPI() |
|
|
|
|
|
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) |
|
|
|
|
|
output = generator(prompt)[0]["generated_text"] |
|
return JSONResponse(content={"response": output}) |
|
|
|
|