ruifma/chatyuan repository overview

作者：苏洋
链接：https://www.zhihu.com/question/583024028/answer/2922406418
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

FROM nvcr.io/nvidia/pytorch:23.01-py3

RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

RUN git clone https://github.com/huggingface/transformers.git --depth=1 --branch=main && \
    pip install ./transformers && \
    rm -rf transformers && \
    pip install sentencepiece

RUN cat > /get-models.py <<EOF
from transformers import T5Tokenizer, T5ForConditionalGeneration
tokenizer = T5Tokenizer.from_pretrained("ClueAI/ChatYuan-large-v1")
model = T5ForConditionalGeneration.from_pretrained("ClueAI/ChatYuan-large-v1")
EOF

RUN python /get-models.py && \
    rm -rf /get-models.py

docker run --gpus all --ipc=host --ulimit memlock=-1  -it --rm chatyuan python
or 
docker run --ipc=host --ulimit memlock=-1  -it --rm chatyuan python

作者：苏洋
链接：https://www.zhihu.com/question/583024028/answer/2922406418
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch

model_name = "ClueAI/ChatYuan-large-v1"
tokenizer = T5Tokenizer.from_pretrained(model_name)
device = torch.device('cuda')
model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)

def trim(text):
  return text.replace("\\n", "\n").replace("\\t", "\t")

def answer(text, top_p=1, temperature=0.7):
  text = trim(text)
  encoding = tokenizer(text=[text], truncation=True, padding=True, max_length=768, return_tensors="pt").to(device) 
  out = model.generate(**encoding, return_dict_in_generate=True, output_scores=False, max_new_tokens=512, do_sample=True, top_p=top_p, temperature=temperature, no_repeat_ngram_size=3)
  out_text = tokenizer.batch_decode(out["sequences"], skip_special_tokens=True)
  return trim(out_text[0])

def hey(prompt):
    prompt = "用户：" + prompt + "\n小元："
    result = answer(prompt)
    print(f"{prompt}{result}")
    return

def classify(prompt, options):
    prompt = "分类任务：\n" + prompt + "\n选项：" + options + "\n答案："
    result = answer(prompt)
    print(f"{prompt}{result}")
    return

def correct(prompt):
    prompt = "文本纠错：\n" + prompt + "\n答案："
    result = answer(prompt)
    print(f"{prompt}{result}\n")
    return

context = ""
def chat(prompt):
    global context
    prompt = "用户：" + prompt + "\n小元："
    context = context + "\n" + prompt
    result = answer(context, top_p=0.9)
    print(f"{prompt}{result}\n")
    return