Update README.md

README.md

@@ -8,14 +8,17 @@ pipeline_tag: text-generation
 library_name: transformers
 tags:
 - chat
+datasets:
+- simplescaling/s1K-1.1_tokenized
 ---
 
 ## Quickstart
 
-Here provides a code snippet with `apply_chat_template` to show you how to load the tokenizer and model and how to generate contents.
+Here provides a code snippet with `apply_chat_template` to show you how to load the tokenizer and model and how to generate contents. 
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import TextStreamer
 
 model_name = "SurgeGlobal/s1-1.1-Qwen-2.5-1.5B-Instruct"
 
@@ -26,7 +29,9 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-prompt = "Give me a short introduction to large language model."
+streamer = TextStreamer(tokenizer, skip_prompt=True)
+
+prompt = "How would a typical person answer each of the following questions about causation? A machine is set up in such a way that it will short circuit if both the black wire and the red wire touch the battery at the same time. The machine will not short circuit if just one of these wires touches the battery. The black wire is designated as the one that is supposed to touch the battery, while the red wire is supposed to remain in some other part of the machine. One day, the black wire and the red wire both end up touching the battery at the same time. There is a short circuit. Did the black wire cause the short circuit? Options: - Yes - No"
 messages = [
     {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
     {"role": "user", "content": prompt}
@@ -40,13 +45,12 @@ model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
 generated_ids = model.generate(
     **model_inputs,
-    max_new_tokens=512
+    max_new_tokens=8000,
+    streamer=streamer
 )
 generated_ids = [
     output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
 ]
 
 response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-
-print(response)
 ```