把Llama2封装为API服务并做一个互动网页

原创

已于 2024-02-19 23:15:16 修改 · 1k 阅读

12 ·

CC 4.0 BY-SA版权

文章标签：

#python #llama #streamlit

于 2024-02-19 23:14:10 首次发布

本文介绍了如何将Facebook的Llama2模型转换为API服务，使用FastAPI和uvicorn搭建服务器，同时实现了一个类似OpenAI的聊天界面，通过Streamlit集成到一个可交互的对话框中。

最近按照官方例子，把Llama2跑起来了测试通了，但是想封装成api服务，耗费了一些些力气

参考：https://github.com/facebookresearch/llama/pull/147/files

1. 准备的前提如下

按照官方如下命令，可以运行成功

torchrun --nproc_per_node 1 example_chat_completion.py \
   --ckpt_dir llama-2-7b-chat/ \
   --tokenizer_path tokenizer.model \
   --max_seq_len 512 --max_batch_size 6

使用的模型是llama-2-7b-chat

2. 第一步，增加依赖包

fastapi
uvicorn

3. 第二步，增加文件`server.py`在`llama`仓库的根目录下

from typing import Tuple
import os
import sys
import argparse
import torch
import time
import json

from pathlib import Path
from typing import List

from pydantic import BaseModel
from fastapi import FastAPI
import uvicorn
import torch.distributed as dist

from fairscale.nn.model_parallel.initialize import initialize_model_parallel

from llama import ModelArgs, Transformer, Tokenizer, Llama


parser = argparse.ArgumentParser()
parser.add_argument('--ckpt_dir', type=str, default='llama-2-7b-chat')
parser.add_argument('--tokenizer_path', type=str, default='tokenizer.model')
parser.add_argument('--max_seq_len', type=int, default=512)
parser.add_argument('--max_batch_size', type=int, default=6)

os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12345'
os.environ['WORLD_SIZE'] = '1'

app = FastAPI()

def setup_model_parallel() -> Tuple