CodeGeeX2模型安装

liuzhenghua66

已于 2023-11-27 17:29:42 修改

阅读量948

点赞数

分类专栏：环境搭建文章标签：机器学习

于 2023-11-27 17:27:03 首次发布

本文链接：https://blog.csdn.net/liuzhenghua66/article/details/134642165

版权

环境搭建专栏收录该内容

18 篇文章 0 订阅

订阅专栏

简介

CodeGeeX2 是多语言代码生成模型 CodeGeeX (KDD’23) 的第二代模型。不同于一代 CodeGeeX（完全在国产华为昇腾芯片平台训练），CodeGeeX2 是基于 ChatGLM2 架构加入代码预训练实现，得益于 ChatGLM2 的更优性能，CodeGeeX2 在多项指标上取得性能提升（+107% > CodeGeeX；仅60亿参数即超过150亿参数的 StarCoder-15B 近10%）
官方仓库地址：https://github.com/THUDM/CodeGeeX2/tree/main

部署步骤

安装依赖

conda create --name CodeGeeX2 python=3.9
conda activate CodeGeeX2

// #或者手动下载并上传至服务器
git clone git@github.com:THUDM/CodeGeeX2.git
cd CodeGeeX2
conda install -y pytorch==2.0.0 torchvision==0.15.0 torchaudio==2.0.0 pytorch-cuda=11.7 -c pytorch -c nvidia
# 一般来说，不需要安装这个，用显卡驱动自带的就行。
# conda install -y cuda-toolkit=11.7.0 -c nvidia

pip install -r requirements.txt -i https://mirrors.aliyun.com/pypi/simple/

下载模型数据

由于huggingface国内无法访问，我们使用镜像站下载

pip install -U huggingface_hub -i https://mirrors.aliyun.com/pypi/simple/
export HF_ENDPOINT=https://hf-mirror.com
# 下载THUDM/codegeex2-6b模型到本地目录：/mnt/data/CodeGeeX2/model
huggingface-cli download --resume-download --local-dir-use-symlinks False THUDM/codegeex2-6b --cache-dir /mnt/data/CodeGeeX2/model

启动

quick start - 单gpu

from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True, cache_dir="/mnt/data/CodeGeeX2/model")
model = AutoModel.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True, device='cuda', cache_dir="/mnt/data/CodeGeeX2/model")
model = model.eval()

# remember adding a language tag for better performance
prompt = "# language: Python\n# write a bubble sort function\n"
inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(inputs, max_length=256, top_k=1)
response = tokenizer.decode(outputs[0])

print(response)

quick start - 多gpu
demo.gpus为CodeGeeX2官方仓库的demo文件夹下的gps文件，需要修改num_gpus为机器实际的值，可以通过nvidia-smi -L查看GPU个数。

from transformers import AutoTokenizer, AutoModel
from demo.gpus import load_model_on_gpus

def get_model():
    tokenizer = AutoTokenizer.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True, cache_dir="/mnt/data/CodeGeeX2/model")
    # gpus文件在demo文件夹中
    model = load_model_on_gpus("THUDM/codegeex2-6b", num_gpus=4, cache_dir="/mnt/data/CodeGeeX2/model")
    model = model.eval()
    return tokenizer, model

tokenizer, model = get_model()

# remember adding a language tag for better performance
prompt = "# language: Python\n# write a bubble sort function\n"
inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(inputs, max_length=256, top_k=1)
response = tokenizer.decode(outputs[0])

print(response)

quick start - http api

from fastapi import FastAPI, Request
from transformers import AutoTokenizer, AutoModel
from demo.gpus import load_model_on_gpus
import uvicorn, json, datetime
import argparse


#获取选项        
def add_code_generation_args(parser):
    group = parser.add_argument_group(title="CodeGeeX2 DEMO")
    group.add_argument(
        "--model-path",
        type=str,
        default="THUDM/codegeex2-6b",
    )
    group.add_argument(
        "--listen",
        type=str,
        default="127.0.0.1",
    )
    group.add_argument(
        "--port",
        type=int,
        default=7860,
    )
    group.add_argument(
        "--workers",
        type=int,
        default=1,
    )
    group.add_argument(                      
        "--cpu",
        action="store_true",
    )
    group.add_argument(                      
        "--half",
        action="store_true",
    )
    group.add_argument(
        "--quantize",
        type=int,
        default=None,
    )
    group.add_argument(
        "--chatglm-cpp",
        action="store_true",
    )
    return parser

LANGUAGE_TAG = {
    "Abap"         : "* language: Abap",
    "ActionScript" : "// language: ActionScript",
    "Ada"          : "-- language: Ada",
    "Agda"         : "-- language: Agda",
    "ANTLR"        : "// language: ANTLR",
    "AppleScript"  : "-- language: AppleScript",
    "Assembly"     : "; language: Assembly",
    "Augeas"       : "// language: Augeas",
    "AWK"          : "// language: AWK",
    "Basic"        : "' language: Basic",
    "C"            : "// language: C",
    "C#"           : "// language: C#",
    "C++"          : "// language: C++",
    "CMake"        : "# language: CMake",
    "Cobol"        : "// language: Cobol",
    "CSS"          : "/* language: CSS */",
    "CUDA"         : "// language: Cuda",
    "Dart"         : "// language: Dart",
    "Delphi"       : "{language: Delphi}",
    "Dockerfile"   : "# language: Dockerfile",
    "Elixir"       : "# language: Elixir",
    "Erlang"       : f"% language: Erlang",
    "Excel"        : "' language: Excel",
    "F#"           : "// language: F#",
    "Fortran"      : "!language: Fortran",
    "GDScript"     : "# language: GDScript",
    "GLSL"         : "// language: GLSL",
    "Go"           : "// language: Go",
    "Groovy"       : "// language: Groovy",
    "Haskell"      : "-- language: Haskell",
    "HTML"         : "<!--language: HTML-->",
    "Isabelle"     : "(*language: Isabelle*)",
    "Java"         : "// language: Java",
    "JavaScript"   : "// language: JavaScript",
    "Julia"        : "# language: Julia",
    "Kotlin"       : "// language: Kotlin",
    "Lean"         : "-- language: Lean",
    "Lisp"         : "; language: Lisp",
    "Lua"          : "// language: Lua",
    "Markdown"     : "<!--language: Markdown-->",
    "Matlab"       : f"% language: Matlab",
    "Objective-C"  : "// language: Objective-C",
    "Objective-C++": "// language: Objective-C++",
    "Pascal"       : "// language: Pascal",
    "Perl"         : "# language: Perl",
    "PHP"          : "// language: PHP",
    "PowerShell"   : "# language: PowerShell",
    "Prolog"       : f"% language: Prolog",
    "Python"       : "# language: Python",
    "R"            : "# language: R",
    "Racket"       : "; language: Racket",
    "RMarkdown"    : "# language: RMarkdown",
    "Ruby"         : "# language: Ruby",
    "Rust"         : "// language: Rust",
    "Scala"        : "// language: Scala",
    "Scheme"       : "; language: Scheme",
    "Shell"        : "# language: Shell",
    "Solidity"     : "// language: Solidity",
    "SPARQL"       : "# language: SPARQL",
    "SQL"          : "-- language: SQL",
    "Swift"        : "// language: swift",
    "TeX"          : f"% language: TeX",
    "Thrift"       : "/* language: Thrift */",
    "TypeScript"   : "// language: TypeScript",
    "Vue"          : "<!--language: Vue-->",
    "Verilog"      : "// language: Verilog",
    "Visual Basic" : "' language: Visual Basic",
}

app = FastAPI()

@app.post("/")
async def create_item(request: Request):
    global model, tokenizer
    json_post_raw = await request.json()
    json_post = json.dumps(json_post_raw)
    json_post_list = json.loads(json_post)
    lang = json_post_list.get('lang')
    prompt = json_post_list.get('prompt')
    max_length = json_post_list.get('max_length', 128)
    top_p = json_post_list.get('top_p', 0.95)
    temperature = json_post_list.get('temperature', 0.2)
    top_k = json_post_list.get('top_k', 1)
    if lang != "None":
        prompt = LANGUAGE_TAG[lang] + "\n" + prompt
    inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
    response = model.generate(inputs,
                              max_length=max_length,
                              do_sample=temperature > 0,
                              top_p=top_p,
                              top_k=top_k,
                              temperature=temperature)
    now = datetime.datetime.now()
    time = now.strftime("%Y-%m-%d %H:%M:%S")
    answer = {
        "response": response,
        "lang": lang,
        "status": 200,
        "time": time
    }
    log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"'
    print(log)

    return answer

def get_model():
    tokenizer = AutoTokenizer.from_pretrained("THUDM/codegeex2-6b", trust_remote_code=True, cache_dir="/mnt/data/CodeGeeX2/model")
    # gpus文件在demo文件夹中
    model = load_model_on_gpus("THUDM/codegeex2-6b", num_gpus=3, cache_dir="/mnt/data/CodeGeeX2/model")
    model = model.eval()
    return tokenizer, model

if __name__ == '__main__':    
    parser = argparse.ArgumentParser()
    parser = add_code_generation_args(parser)
    args, _ = parser.parse_known_args()
    tokenizer, model = get_model()
    uvicorn.run(app, host=args.listen, port=args.port, workers=args.workers)

启动命令

python web_quick_start.py --listen '0.0.0.0' --port 7860 --workders 100

# 另一个终端中调用
curl -X POST "http://127.0.0.1:7860" \
    -H 'Content-Type: application/json' \
    -d '{"lang": "Python", "prompt": "# Write a quick sort function"}'

参考资料

校验torch安装结果

import torch
print(torch.cuda.is_available())
# 返回gpu数量；
print(torch.cuda.device_count())
# 返回gpu名字，设备索引默认从0开始；
print(torch.cuda.get_device_name(0))
print(torch.cuda.current_device())

# 查看nvidia cuda toolkit版本
nvcc -V

版本查询
cuda tookit版本：https://anaconda.org/nvidia/cuda-toolkit
pytorch与cuda对应版本：https://pytorch.org/get-started/previous-versions/

liuzhenghua66

关注

0
点赞
踩
4

收藏

觉得还不错? 一键收藏
0
评论
CodeGeeX2模型安装

pytorch与cuda对应版本：https://pytorch.org/get-started/previous-versions/demo.gpus为CodeGeeX2官方仓库的demo文件夹下的gps文件，需要修改num_gpus为机器实际的值，可以通过。cuda tookit版本：https://anaconda.org/nvidia/cuda-toolkit。由于huggingface国内无法访问，我们使用镜像站下载。
复制链接

扫一扫

专栏目录