From 8e74007ef8cf8852ad7a9b5b88fd21c537dc7622 Mon Sep 17 00:00:00 2001
From: MengAiDev <3463526515@qq.com>
Date: Sun, 6 Apr 2025 09:36:19 +0000
Subject: [PATCH] add tests

---
 README.md                 | 12 ++++++++++++
 tests/test_fusedbitnet.py | 12 ++++++++++++
 tests/test_generate.py    |  7 +++++++
 tests/test_layernorm.py   | 16 ++++++++++++++++
 tests/test_utils.py       | 11 +++++++++++
 5 files changed, 58 insertions(+)
 create mode 100644 tests/test_fusedbitnet.py
 create mode 100644 tests/test_generate.py
 create mode 100644 tests/test_layernorm.py
 create mode 100644 tests/test_utils.py

diff --git a/README.md b/README.md
index 3cc7735..a8a957a 100644
--- a/README.md
+++ b/README.md
@@ -114,7 +114,19 @@ outputs = model.generate(input_ids, max_length=32,  do_sample=True, top_p=0.4, t
 print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
 ```
 
+# Running Tests
 
+To run the tests, ensure you have `pytest` installed:
+
+```sh
+pip install pytest
+```
+
+Run the tests using:
+
+```sh
+pytest /workspaces/matmulfreellm/tests
+```
 
 # Citation
 If you use this repo in your work, please cite our preprint:
diff --git a/tests/test_fusedbitnet.py b/tests/test_fusedbitnet.py
new file mode 100644
index 0000000..8359404
--- /dev/null
+++ b/tests/test_fusedbitnet.py
@@ -0,0 +1,12 @@
+import torch
+from mmfreelm.ops.fusedbitnet import activation_quant, weight_quant
+
+def test_activation_quant():
+    x = torch.tensor([1.0, 2.0, 3.0])
+    result = activation_quant(x)
+    assert result is not None  # Add more specific assertions based on expected behavior
+
+def test_weight_quant():
+    w = torch.tensor([1.0, 2.0, 3.0])
+    result = weight_quant(w)
+    assert result is not None  # Add more specific assertions based on expected behavior
diff --git a/tests/test_generate.py b/tests/test_generate.py
new file mode 100644
index 0000000..2261ec7
--- /dev/null
+++ b/tests/test_generate.py
@@ -0,0 +1,7 @@
+from transformers import AutoTokenizer
+
+def test_generate_script():
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")  # Replace with a valid model name
+    input_prompt = "Test prompt"
+    input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids
+    assert input_ids is not None
diff --git a/tests/test_layernorm.py b/tests/test_layernorm.py
new file mode 100644
index 0000000..d2aeaf4
--- /dev/null
+++ b/tests/test_layernorm.py
@@ -0,0 +1,16 @@
+import torch
+from mmfreelm.modules.layernorm import layer_norm_ref, rms_norm_ref
+
+def test_layer_norm_ref():
+    x = torch.randn(10, 10)
+    weight = torch.ones(10)
+    bias = torch.zeros(10)
+    result = layer_norm_ref(x, weight, bias)
+    assert result is not None  # Add more specific assertions based on expected behavior
+
+def test_rms_norm_ref():
+    x = torch.randn(10, 10)
+    weight = torch.ones(10)
+    bias = torch.zeros(10)
+    result = rms_norm_ref(x, weight, bias)
+    assert result is not None  # Add more specific assertions based on expected behavior
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..e9ea117
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,11 @@
+import torch
+from mmfreelm.utils import contiguous
+
+def test_contiguous():
+    @contiguous
+    def dummy_fn(ctx, x):
+        return x
+
+    x = torch.randn(10, 10).t()  # Non-contiguous tensor
+    result = dummy_fn(None, x)
+    assert result.is_contiguous()