fastapi_t5 / main.py
real-jiakai's picture
Update main.py
9a1d2d6 verified
raw
history blame
1.58 kB
# 从 FastAPI 库中导入 FastAPI 类,用于创建应用实例
from fastapi import FastAPI
# 从 FastAPI 的 staticfiles 模块中导入 StaticFiles 类,用于处理静态文件
from fastapi.staticfiles import StaticFiles
# 从 FastAPI 的 responses 模块中导入 FileResponse 类,用于返回文件响应
from fastapi.responses import FileResponse
# 从 transformers 库中导入 pipeline 函数,用于加载预训练模型管道
from transformers import pipeline
# 创建 FastAPI 应用实例
app = FastAPI()
# 使用 transformers 的 pipeline 函数加载 Flan-T5 小模型,用于文本生成任务
pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")
# 定义一个 GET 请求的路由,路径为 /infer_t5,用于处理文本生成请求
@app.get("/infer_t5")
# 定义处理 /infer_t5 路由的函数,接收一个输入参数
def t5(input):
# 使用加载的 Flan-T5 模型管道生成输出
output = pipe_flan(input)
# 返回生成的文本结果,格式为 JSON
return {"output": output[0]["generated_text"]}
# 挂载静态文件目录,将根路径 "/" 指向 "static" 目录,并启用 HTML 模式
app.mount("/", StaticFiles(directory="static", html=True), name="static")
# 定义一个 GET 请求的路由,路径为根路径 "/",用于返回首页
@app.get("/")
# 定义处理根路径的函数,返回一个文件响应,类型为 HTML
def index() -> FileResponse:
# 返回位于 /app/static/index.html 的 HTML 文件
return FileResponse(path="/app/static/index.html", media_type="text/html")