本次案例 vllm
结合 transformers
的 AutoTokenizer
来加载本地模型进行推理。支持异步流式返回生成的文本
优化后的代码示例:
import logging
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer
import uvicorn
import asyncio
# 定义请求的数据模型
class PromptRequest(BaseModel):
prompt: str
# 初始化 FastAPI 应用
app = FastAPI()
# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# 并发请求限制
MAX_CONCURRENT_REQUESTS = 10
# 等待队列最大长度
MAX_QUEUE_SIZE = 20
# 创建信号量和队列
semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
request_queue = asyncio.Queue(MAX_QUEUE_SIZE)
# 本地模型路径
model_path = "path_to_your_model" # 指定本地模型文件夹路径