Update README.md

README.md

@@ -34,22 +34,50 @@ Below is a simple example demonstrating how to use the model for solving a mathe
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 # Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("haijian06/Yi-1.5-6B-Chat-Math")
-model = AutoModelForCausalLM.from_pretrained("haijian06/Yi-1.5-6B-Chat-Math")
+model = AutoModelForCausalLM.from_pretrained("haijian06/Yi-1.5-6B-Chat-Math", torch_dtype=torch.float16, device_map="auto")
 
-# Input mathematical problem
-input_text = "Solve the equation x^2 - 5x + 6 = 0"
-inputs = tokenizer(input_text, return_tensors="pt")
+input_text = "Solve the equation x^2 - 5x + 6 = 0 Let's solve this step-by-step:"
 
-# Generate the solution
-outputs = model.generate(**inputs)
-answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+
+with torch.no_grad():
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=200,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.95,
+    )
 
+answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
 print(answer)
 ```
+```
+Solve the equation x^2  - 5 x + 6  = 0  Let's solve this step-by-step:
+
+Step 1: Factor the equation
+The equation can be factored as follows:
+
+x^2 - 5x + 6 = 0
+(x - 2)(x - 3) = 0
 
+Step 2: Apply the zero product property
+If the product of two numbers is zero, then at least one of the numbers must be zero.
+
+So, either (x - 2) = 0 or (x - 3) = 0
+
+Step 3: Solve for x
+If (x - 2) = 0, then x = 2
+If (x - 3) = 0, then x = 3
+
+So, the solutions are x = 2 and x = 3.
+
+Answer: 2, 3
+```
 ## Contributing
 
 Contributions are welcome! Whether you have suggestions for improvements, bug reports, or want to contribute code, feel free to open an issue or submit a pull request on GitHub.