Update README.md
Browse files
README.md
CHANGED
|
@@ -34,7 +34,7 @@ import torch
|
|
| 34 |
from transformers import AutoTokenizer, AutoModelForCausalLM
|
| 35 |
tokenizer = AutoTokenizer.from_pretrained('CofeAI/Tele-FLM', trust_remote_code=True)
|
| 36 |
model = AutoModelForCausalLM.from_pretrained('CofeAI/Tele-FLM', torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map="auto", trust_remote_code=True)
|
| 37 |
-
inputs = tokenizer('
|
| 38 |
generated = model.generate(**inputs, max_new_tokens=128, repetition_penalty=1.03)
|
| 39 |
print(tokenizer.decode(generated.cpu()[0], skip_special_tokens=True))
|
| 40 |
```
|
|
|
|
| 34 |
from transformers import AutoTokenizer, AutoModelForCausalLM
|
| 35 |
tokenizer = AutoTokenizer.from_pretrained('CofeAI/Tele-FLM', trust_remote_code=True)
|
| 36 |
model = AutoModelForCausalLM.from_pretrained('CofeAI/Tele-FLM', torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map="auto", trust_remote_code=True)
|
| 37 |
+
inputs = tokenizer('Beijing is the capital of China.', return_tensors='pt').to(model.device)
|
| 38 |
generated = model.generate(**inputs, max_new_tokens=128, repetition_penalty=1.03)
|
| 39 |
print(tokenizer.decode(generated.cpu()[0], skip_special_tokens=True))
|
| 40 |
```
|