Update README.md

README.md

@@ -47,6 +47,8 @@ extra_gated_button_content: Submit
 base_model: meta-llama/Llama-3.2-3B-Instruct
 metrics:
 - accuracy
+datasets:
+- HenryShan/Gemini-MMLU-CoT
 ---
 # MathLlama 3.2 - Enhanced Mathematical Reasoning Model
 
@@ -57,21 +59,15 @@ MathLlama 3.2 is a fine-tuned version of Meta's Llama-3.2-3B-Instruct model, spe
 ## Key Improvements
 
 ### Mathematical Reasoning Enhancement
-- **10% improvement on abstract_algebra in MMLU** compared to the base Llama-3.2-3B-Instruct model
+- **12% improvement on abstract_algebra in MMLU** compared to the base Llama-3.2-3B-Instruct model
 - Enhanced capability in complex mathematical reasoning tasks
 - Improved performance across various mathematical domains including algebra, calculus, and abstract mathematical concepts
 
 ### Training Methodology
-- **Synthetic Dataset**: Utilized a self-created synthetic dataset consisting of **1,000 advanced math problems**
+- **Synthetic Dataset**: Utilized the [Gemini-MMLU-CoT](https://huggingface.co/datasets/HenryShan/Gemini-MMLU-CoT) dataset consisting of **7,000 advanced math problems**
 - **Chain-of-Thought Training**: Each training example includes detailed step-by-step reasoning processes
-- **Data Quality**: Carefully curated advanced mathematical problems covering:
-  - Abstract algebra concepts (groups, rings, fields)
-  - Advanced calculus problems
-  - Complex mathematical proofs
-  - Multi-step mathematical reasoning tasks
 
 ## Model Architecture
-
 - **Base Model**: Meta Llama-3.2-3B-Instruct
 - **Parameters**: 3 billion parameters
 - **Context Length**: 128k tokens
@@ -81,39 +77,32 @@ MathLlama 3.2 is a fine-tuned version of Meta's Llama-3.2-3B-Instruct model, spe
 
 ### Basic Usage
 ```python
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-tokenizer = AutoTokenizer.from_pretrained("path/to/MathLlama-3.2")
-model = AutoModelForCausalLM.from_pretrained("path/to/MathLlama-3.2")
 
-# Example mathematical reasoning
-prompt = "Solve the following algebra problem step by step: Find x where 2x + 3 = 7"
-inputs = tokenizer(prompt, return_tensors="pt")
-
-outputs = model.generate(**inputs, max_length=512, temperature=0.7)
-response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-print(response)
-```
-
-### Advanced Mathematical Reasoning
-```python
-# For complex mathematical problems
-complex_prompt = """
-Solve this advanced algebra problem using step-by-step reasoning:
-
-Let G be a group and H be a subgroup of G. If [G:H] = n and for every g in G,
-the order of g divides some fixed integer m, prove that the order of G divides n! * m^n.
+# Load model directly
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-Show your step-by-step reasoning:
-"""
+tokenizer = AutoTokenizer.from_pretrained("HenryShan/MathLlama3.2")
+model = AutoModelForCausalLM.from_pretrained("HenryShan/MathLlama3.2")
+messages = [
+    {"role": "user", "content": "Find the degree for the given field extension Q(sqrt(2), sqrt(3), sqrt(18)) over Q. Answer Choices: A: "0", B: "4", C: "2", D: "6"},
+]
+inputs = tokenizer.apply_chat_template(
+	messages,
+	add_generation_prompt=True,
+	tokenize=True,
+	return_dict=True,
+	return_tensors="pt",
+).to(model.device)
+
+outputs = model.generate(**inputs, max_new_tokens=40)
+print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))
 ```
 
 ## Training Details
 
 ### Dataset Creation
-- **Self-Created Synthetic Data**: 1,000 carefully designed advanced mathematical problems
-- **Chain-of-Thought Format**: Each problem includes:
+- **Synthetic Dataset**: 7,000 carefully designed advanced mathematical problems
+- **Chain-of-Thought Format**: Each row includes:
   1. Clear problem statement
   2. Step-by-step reasoning process
   3. Final answer with justification
@@ -122,7 +111,7 @@ Show your step-by-step reasoning:
 - **Learning Rate**: Optimized for mathematical reasoning tasks
 - **Batch Size**: Configured for stable training with mathematical data
 - **Training Steps**: Sufficient iterations to achieve mathematical reasoning improvements
-- **Hardware**: Trained on modern GPU infrastructure
+- **Hardware**: Trained on an Apple M4 Max Computer using [MLX](https://github.com/ml-explore/mlx)
 
 ## Applications