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="auto")
tokenizer = AutoTokenizer.from_pretrained(model_path)
@app.post("/chat/")
async def generate_response(query: Query):
inputs = f"[INST] {query.text.strip()} [/INST]"
input_ids = tokenizer(inputs, return_tensors="pt").input_ids.to(device)
generate_ids = model.generate(
input_ids,
max_new_tokens=500,
do_sample=True,
top_p=0.85,
temperature=1.0,
repetition_penalty=1.,
eos_token_id=2,
bos_token_id=1,
pad_token_id=0)
output = tokenizer.batch_decode(generate_ids)[0]
return {"result": output}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=6006)
客户端代码
import requests
url = "https://xxxxxxxxxxxx/chat/"
# 使用新的输入格式,包裹用户输入
query = {"text": "[INST] introduce china[/INST]"} # 修改为使用[INST]标签
response = requests.post(url, json=query)
if response.status_code == 200:
result = response.json()
print("chat:", result["result"])
else:
print("Error:", response.status_code, response.text)