Fastapi部署大语言模型llama2

vpcyh

已于 2024-01-24 10:30:54 修改

阅读量647

点赞数 11

文章标签： python

于 2024-01-24 10:30:01 首次发布

本文链接：https://blog.csdn.net/qq_43587354/article/details/135815636

版权

本文介绍了如何使用FastAPI框架部署一个基于HuggingFaceLlama模型的文本生成聊天服务器，以及客户端如何发送请求获取响应。

摘要由CSDN通过智能技术生成

Fastapi部署llama

服务端代码

import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, LlamaForCausalLM
import torch

app = FastAPI()

class Query(BaseModel):
    text: str

device = torch.device("cuda:0")

model_path = 'llama-2-7b-chat-hf'
model = LlamaForCausalLM.from_pretrained(model_path, device_map=