Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3tokenizer = AutoTokenizer("google/gemma-7b")
4model = AutoModelForCausalLM.from_pretrained(
5 "merve/gemma-7b-8bit",
6 device_map='auto'
7)
8
9input_text = "Write me a poem about Machine Learning."
10input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
11
12outputs = model.generate(**input_ids)
13print(tokenizer.decode(outputs[0]))