https://github.com/jzhang38/TinyLlama
TinyLlama 项目旨在预训练一个基于 3 万亿 tokens 的 1.1B Llama 模型。通过适当的优化,我们仅需 16 张 A100-40G GPU,就能在“短短”90 天内完成这一目标 🚀🚀。训练已于 2023 年 9 月 1 日启动。
我们采用了与 Llama 2 完全相同的架构和分词器。这意味着 TinyLlama 可以直接应用于许多基于 Llama 构建的开源项目。此外,TinyLlama 仅有 1.1B 参数,非常小巧。这种小巧的特性使其能够满足众多对计算和内存资源有严格限制的应用场景。
这是在 [TinyLlama/TinyLlama-1.1B-intermediate-step-715k-1.5T] 基础上微调得到的对话模型。 所使用的数据集是 [OpenAssistant/oasst_top1_2023-08-25],并遵循 chatml 格式。
您需要安装 transformers>=4.31 更多信息请查看 TinyLlama 的 GitHub 页面。
from openmind import AutoTokenizer
import openmind
import torch
model = "LF_AICC/TinyLlama-1.1B-Chat-v0.4"
tokenizer = AutoTokenizer.from_pretrained(model)
pipeline = openmind.pipeline(
"text-generation",
model=model,
torch_dtype=torch.float16,
device_map="auto",
)
CHAT_EOS_TOKEN_ID = 32002
prompt = "How to get in a good university?"
formatted_prompt = (
f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
)
sequences = pipeline(
formatted_prompt,
do_sample=True,
top_k=50,
top_p = 0.9,
num_return_sequences=1,
repetition_penalty=1.1,
max_new_tokens=1024,
eos_token_id=CHAT_EOS_TOKEN_ID,
)
for seq in sequences:
print(f"Result: {seq['generated_text']}")