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 |
```
|