Update README.md

README.md

@@ -22,14 +22,23 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 tokenizer_name = "Google/gemma-2-2b-it"
 model_name="lmassaron/gemma-2-2b-it-grpo-gsm8k"
-tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", use_cache=True)
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_name,
+                                          trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_name,
+                                             device_map="auto",
+                                             use_cache=True)
 
 FORMAT = """<reasoning>\n</reasoning>\n<answer>\n</answer>\n"""
 
 question = "Which is bigger? 9.11 or 9.9?"
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, do_sample=False, batch_size=1)
-output = generator([{"role": "user", "content": FORMAT + question}], max_new_tokens=256, return_full_text=False)[0]
+generator = pipeline("text-generation",
+                      model=model,
+                      tokenizer=tokenizer,
+                      do_sample=False,
+                      batch_size=1)
+output = generator([{"role": "user", "content": FORMAT + question}],
+                    max_new_tokens=256,
+                    return_full_text=False)[0]
 print(output["generated_text"])
 ```