NLP tasks | |
First, let's set up the environment: | |
pip install -q transformers accelerate | |
Next, let's load the model with the appropriate pipeline ("text-generation"): | |
thon | |
from transformers import pipeline, AutoTokenizer | |
import torch | |
torch.manual_seed(0) # doctest: +IGNORE_RESULT | |
model = "tiiuae/falcon-7b-instruct" | |
tokenizer = AutoTokenizer.from_pretrained(model) | |
pipe = pipeline( | |
"text-generation", | |
model=model, | |
tokenizer=tokenizer, | |
torch_dtype=torch.bfloat16, | |
device_map="auto", | |
) | |
Note that Falcon models were trained using the bfloat16 datatype, so we recommend you use the same. |