12 lines
434 B
Python
12 lines
434 B
Python
|
# pip install accelerate
|
||
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||
|
|
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it")
|
||
|
|
model = AutoModelForCausalLM.from_pretrained("google/gemma-7b-it", device_map="auto")
|
||
|
|
|
||
|
|
input_text = "Write me a poem about Machine Learning."
|
||
|
|
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
|
||
|
|
|
||
|
|
outputs = model.generate(**input_ids)
|
||
|
|
print(tokenizer.decode(outputs[0]))
|