零一科技Yi-34B Chat大模型环境搭建&推理

本文链接：https://blog.csdn.net/zzq1989_/article/details/135597181

本文介绍了如何在Python中安装并使用零一科技的Yi-34B大模型进行对话生成和问题解答，包括环境配置、代码示例以及效果演示。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

引子

　　国产大模型此起彼伏，各种刷榜。作为没有能力训练的我们，只能跟着大佬开源的模型尝试下效果，零一科技23年底发布的Yi大模型，之前就有尝试，现汇总总结下，OK，我们开始吧。

一、环境安装

下载模型

https://huggingface.co/01-ai/Yi-34B-Chat/tree/main

conda activate chatglm

git clone GitHub - 01-ai/Yi: A series of large language models trained from scratch by developers @01-ai

cd /opt/tmp/zzq/Yi/

pip install -r requirements.txt

二、推理代码运行

python chat_yi_34B_demo.py

三、效果展示

1、Q：

A：

2、Q：

A：

四、代码

from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
model_path = 'Yi-34B-Chat'
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
# Since transformers 4.35.0, the GPT-Q/AWQ model can be loaded using AutoModelForCausalLM.
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map="auto",
    torch_dtype='auto'
).eval()
# 对话历史
messages = [{"role": "user", "content": "请扮演一个AI助手角色，你的名字叫小艺"},
            {"role": "assistant", "content": "好的，我记住了。"},
            # {"role": "user", "content": "小艺，吕布是三国人物么？请按照如下格式输出，\
            # 如果是，请回答{'ans':'yes'},如果不是，请回答{'ans':'no'}，请注意不要输出其他信息"}]
            # {"role": "user", "content": "小艺，写一篇赞美雪景的文章，内容包含南方的雪和北方的雪的特色，\
            # 文章字数要求1500字，文章字数要求1500字，文章字数要求1500字"}]
            {"role": "user", "content": "小艺，鸡兔同笼，头共20个，足共62只，求鸡与兔各有多少只？"}]
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
input_ids = tokenizer.apply_chat_template(conversation=messages, tokenize=True, add_generation_prompt=True,
                                          return_tensors='pt')
generation_kwargs = dict(input_ids=input_ids.to('cuda'), streamer=streamer, temperature=0.95, top_p=0.8)
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
for new_text in streamer:
    new_text = new_text.replace('<|im_end|>', '\n')
    print(new_text, end="", flush=True)