Views
No views yet
1from transformers import AutoTokenizer
2import transformers
3import torch
4model = "PY007/TinyLlama-1.1B-intermediate-step-240k-503b"
5tokenizer = AutoTokenizer.from_pretrained(model)
6pipeline = transformers.pipeline(
7 "text-generation",
8 model=model,
9 torch_dtype=torch.float16,
10 device_map="auto",
11)
12
13sequences = pipeline(
14 'The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens. With some proper optimization, we can achieve this within a span of "just" 90 days using 16 A100-40G GPUs 🚀🚀. The training has started on 2023-09-01.',
15 do_sample=True,
16 top_k=10,
17 num_return_sequences=1,
18 repetition_penalty=1.5,
19 eos_token_id=tokenizer.eos_token_id,
20 max_length=500,
21)
22for seq in sequences:
23 print(f"Result: {seq['generated_text']}")