baichuan-7B fastapi接口服务
文章目录
前言
提示:这里可以添加本文要记录的大概内容:
例如:随着人工智能的不断发展,机器学习这门技术也越来越重要,很多人都开启了学习机器学习,本文就介绍了机器学习的基础内容。
提示:以下是本篇文章正文内容,下面案例可供参考
一、环境安装
1 git clone https://github.com/baichuan-inc/baichuan-7B
2 pip install -r requirements.txt
二、baichuan-7B fastapi接口服务
1.baichuan-7B fastapi服务端代码
代码如下(示例):
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer
app = FastAPI()
tokenizer = AutoTokenizer.from_pretrained("baichuan-inc/baichuan-7B", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("baichuan-inc/baichuan-7B", device_map="auto", trust_remote_code=True)
class TextGenerationInput(BaseModel):
text: str
class TextGenerationOutput(BaseModel):
generated_text: str
@app.post("/generate", response_model=TextGenerationOutput)
async def generate_text(input_data: TextGenerationInput):
inputs = tokenizer(input_data.text, return_tensors='pt')
inputs = inputs.to('cuda:0')
pred = model.generate(**inputs, max_new_tokens=64, repetition_penalty=1.1)
generated_text = tokenizer.decode(pred.cpu()[0], skip_special_tokens=True)
return TextGenerationOutput(generated_text=generated_text)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)
2.接口调用代码
代码如下(示例):
import requests
url = "http://127.0.0.1:8000/generate"
data = {
"text": "登鹳雀楼->王之涣\n夜雨寄北->"
}
response = requests.post(url, json=data)
response_data = response.json()
print("Generated text:", response_data["generated_text"])
该处使用的url网络请求的数据。
2.postman curl代码
代码如下(示例):
curl --location --request POST 'http://127.0.0.1:8000/generate' \
--header 'Content-Type: application/json' \
--data-raw '{
"text": "登鹳雀楼->王之涣\n夜雨寄北->"
}'
该处使用的url网络请求的数据。
总结
以上就是今天要讲的内容
需要加大语言模型垂直领域训练交流讨论群的,请私信我