Ouro-1.4B.py 391 Bytes
Newer Older
zzg_666's avatar
zzg_666 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from transformers import pipeline
import torch
pipe = pipeline(
    "text-generation", 
    model="ByteDance/Ouro-1.4B", 
    trust_remote_code=True,
    device_map="auto",  # 自动分配模型到所有可用GPU
    torch_dtype=torch.float16  # 使用半精度减少显存占用
)

messages = [
    {"role": "user", "content": "介绍下你自己"},
]
result = pipe(messages)
print(result)