当前位置: 首页 > article >正文

【AIGC】AIGC程序记录

一、模型相关

二、程序相关

1、GPU使用

方式一:

os.environ["CUDA_VISIBLE_DEVICES"] = "3,4,5"

方式二:

2、是否加载GPU

 -- GPU加载:
 model = AutoModelForCausalLM.from_pretrained(
        "baichuan-inc/Baichuan-13B-Chat",
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True,
    )
    
 -- 先CPU,再GPU加载:
 model = AutoModelForCausalLM.from_pretrained(
        "baichuan-inc/Baichuan-13B-Chat",
        torch_dtype=torch.float16,
        trust_remote_code=True,
    )
 model = model.cuda()
 
 -- 量化:
 model = model.quantize(4).cuda()

三、服务相关

1、chat类-websocket服务

from fastapi import FastAPI, Request, WebSocket
from transformers import AutoTokenizer, AutoModel
import uvicorn, json, datetime
import torch
import os
import asyncio

""" BaiChan13B Example """

#配置GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "3,4,5"
app = FastAPI(debug=False)

@app.websocket("/chat")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    messages = []
    while True:
        prompt = await websocket.receive_text()
        messages.append(
            {"role": "user", "content": prompt_text}
        )
        for response in model.chat(tokenizer, messages, stream=True):
            result = {
                "res":{
                    "text": response,
                    "end": False
                }
            }
            result = json.dumps(result, ensure_ascii=False)
            await websocket.send_text(result)
            await asyncio.sleep(0.001)
        result = {
            "res": {
                "text": response,
                "end": True
            }
        }
        result = json.dumps(result, ensure_ascii=False)
        await websocket.send_text(result)
        await asyncio.sleep(0.01)

        messages.append(
            {"role": "assistant", "content": response}
        )


if __name__ == "__main__":

    model = AutoModelForCausalLM.from_pretrained(
        "baichuan-inc/Baichuan-13B-Chat",
        torch_dtype=torch.float16,
        device_map="auto",
        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
    )

    model.eval()
    uvicorn.run(app, host="0.0.0.0", port=8000, ws_ping_timeout=5)
        

注意

如果websocket返回结果比较慢,加入sleep,如下:

await websocket.send_text(result)
await asyncio.sleep(0.001)

http://www.kler.cn/a/37511.html

相关文章:

  • 为了实现上网自由,我做了一个多功能串口服务器
  • ElasticSearch入门教程--集群搭建和版本比较
  • 100种思维模型之路径依赖思维模型-98
  • 人工智能(pytorch)搭建模型18-含有注意力机制的CoAtNet模型的搭建,加载数据进行模型训练
  • nginx路由
  • 计算机网络 day7 扫描IP脚本 - 路由器 - ping某网址的过程
  • git下载源码及环境搭建下载源码之后端(一)
  • 一道北大强基题背后的故事(七)——特征根公式的来龙去脉
  • Spring - 更简单的获取 Bean 对象
  • Android 四大布局使用详解
  • transformer 学习
  • Filebeat+ELK (grok、mutate、mutiline、date)详解
  • 数据结构与算法——什么是线性表(线性存储结构)
  • NVIDIA Deep Learning Accelerator (DLA) 支持的几种张量格式
  • 3ds Max 无插件制作燃烧的火焰动画特效
  • 直方图 直方图均衡化
  • VScode 终端无法识别npm以及Missing script: “serve“ 问题
  • SpringBoot第22讲:SpringBoot如何实现接口限流之分布式
  • 如何保证消息的可靠性+延迟队列(TTL+死信队列+延迟队列)
  • 小程序:页面跳转闪屏