计算语言模型计算每秒钟生成的token数量it/s

main() 函数的stream循环中,我们可以计算每秒钟生成的token数量,然后输出 it/s。在流式生成过程中,我们可以使用Python的time模块来计算速度。在测试时,生成速度会受到多个因素的影响,包括设备性能、模型大小、输入文本长度等。

python 复制代码
import os
import torch
import platform
from colorama import Fore, Style
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.utils import GenerationConfig
import time


def init_model():
    print("init model ...")
    model = AutoModelForCausalLM.from_pretrained(
        "baichuan-inc/Baichuan-13B-Chat",
        torch_dtype=torch.float16,
        device_map="cuda",
        trust_remote_code=True
    )

    model.generation_config = GenerationConfig.from_pretrained(
        "baichuan-inc/Baichuan-13B-Chat"
    )
    tokenizer = AutoTokenizer.from_pretrained(
        "baichuan-inc/Baichuan-13B-Chat",
        use_fast=False,
        trust_remote_code=True
    )
    return model, tokenizer


def clear_screen():
    if platform.system() == "Windows":
        os.system("cls")
    else:
        os.system("clear")
    print(Fore.YELLOW + Style.BRIGHT + "欢迎使用百川大模型,输入进行对话,clear 清空历史,CTRL+C 中断生成,stream 开关流式生成,exit 结束。")
    return []


def main(stream=True):
    model, tokenizer = init_model()

    messages = clear_screen()
    while True:
        prompt = input(Fore.GREEN + Style.BRIGHT + "\n用户:" + Style.NORMAL)
        if prompt.strip() == "exit":
            break
        if prompt.strip() == "clear":
            messages = clear_screen()
            continue
        print(Fore.CYAN + Style.BRIGHT + "\nBaichuan:" + Style.NORMAL, end='')
        if prompt.strip() == "stream":
            stream = not stream
            print(Fore.YELLOW + "({}流式生成)\n".format("开启" if stream else "关闭"), end='')
            continue
        messages.append({"role": "user", "content": prompt})
        if stream:
            position = 0
            try:
                start_time = time.time()
                total_tokens = 0
                for response in model.chat(tokenizer, messages, stream=True):
                    print(response[position:], end='', flush=True)
                    position = len(response)
                    total_tokens += len(tokenizer(response, return_tensors='pt')['input_ids'][0])
                    if torch.backends.mps.is_available():
                        torch.mps.empty_cache()
                end_time = time.time()
                elapsed_time = end_time - start_time
                tokens_per_second = total_tokens / elapsed_time
                print(f"\n\n生成速度:{tokens_per_second:.2f} tokens/s")
            except KeyboardInterrupt:
                pass
            print()

        else:
            response = model.chat(tokenizer, messages)
            print(response)
            if torch.backends.mps.is_available():
                torch.mps.empty_cache()
        messages.append({"role": "assistant", "content": response})

    print(Style.RESET_ALL)


if __name__ == "__main__":
    main()
相关推荐
SmartBrain2 小时前
DeerFlow 实践:华为IPD流程的评审智能体设计
人工智能·语言模型·架构
九章云极AladdinEdu5 小时前
超参数自动化调优指南:Optuna vs. Ray Tune 对比评测
运维·人工智能·深度学习·ai·自动化·gpu算力
研梦非凡7 小时前
ICCV 2025|从粗到细:用于高效3D高斯溅射的可学习离散小波变换
人工智能·深度学习·学习·3d
JoannaJuanCV9 小时前
大语言模型基石:Transformer
人工智能·语言模型·transformer
通街市密人有9 小时前
IDF: Iterative Dynamic Filtering Networks for Generalizable Image Denoising
人工智能·深度学习·计算机视觉
大千AI助手9 小时前
TruthfulQA:衡量语言模型真实性的基准
人工智能·语言模型·自然语言处理·llm·模型评估·truthfulqa·事实性基准
智数研析社9 小时前
9120 部 TMDb 高分电影数据集 | 7 列全维度指标 (评分 / 热度 / 剧情)+API 权威源 | 电影趋势分析 / 推荐系统 / NLP 建模用
大数据·人工智能·python·深度学习·数据分析·数据集·数据清洗
七元权11 小时前
论文阅读-Correlate and Excite
论文阅读·深度学习·注意力机制·双目深度估计
张较瘦_11 小时前
[论文阅读] 人工智能 + 软件工程 | 大语言模型驱动的多来源漏洞影响库识别研究解析
论文阅读·人工智能·语言模型
什么都想学的阿超11 小时前
【大语言模型 58】分布式文件系统:训练数据高效存储
人工智能·语言模型·自然语言处理