publicTestApp / main.py
dtyago's picture
Initial commit
7efd925
raw
history blame
No virus
859 Bytes
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from transformers import pipeline
app = FastAPI()
# Add middleware
# Set all origins to wildcard for simplicity, but you should limit this in production
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")
@app.get("/infer_t5")
def t5(input):
output = pipe_flan(input)
return {"output": output[0]["generated_text"]}
app.mount("/", StaticFiles(directory="static", html=True), name="static")
@app.get("/")
def index() -> FileResponse:
return FileResponse(path="/app/static/index.html", media_type="text/html")