add plan

2024-03-22 00:30:07 +08:00
parent 47c0e23250
commit ce9f6b6145
3 changed files with 26 additions and 4 deletions
--- a/02.ChatWithGemma7b
+++ b/02.ChatWithGemma7b
@@ -1,11 +1,10 @@
-# 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")
+model = AutoModelForCausalLM.from_pretrained("google/gemma-7b-it")

 input_text = "Write me a poem about Machine Learning."
-input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
+input_ids = tokenizer(input_text, return_tensors="pt")

 outputs = model.generate(**input_ids)
 print(tokenizer.decode(outputs[0]))