Unit tests for pricer-module

week8/community_contributions/pricer_test/tests/test_lib.py

@@ -1,5 +0,0 @@
-from my_pkg.lib import has_gpu
-
-
-def test_torch_cuda():
-    assert has_gpu()

week8/community_contributions/pricer_test/tests/test_pricer.py

@@ -0,0 +1,84 @@
+import pdb
+from pricer.ci import Pricer
+from unittest.mock import patch, MagicMock
+import torch
+import pytest
+from transformers import BitsAndBytesConfig
+
+BASE_MODEL = "meta-llama/Meta-Llama-3.1-8B"
+PROJECT_NAME = "pricer"
+HF_USER = "ed-donner" # your HF name here! Or use mine if you just want to reproduce my results.
+RUN_NAME = "2024-09-13_13.04.39"
+PROJECT_RUN_NAME = f"{PROJECT_NAME}-{RUN_NAME}"
+REVISION = "e8d637df551603dc86cd7a1598a8f44af4d7ae36"
+FINETUNED_MODEL = f"{HF_USER}/{PROJECT_RUN_NAME}"
+MODEL_DIR = "hf-cache/"
+BASE_DIR = MODEL_DIR + BASE_MODEL
+FINETUNED_DIR = MODEL_DIR + FINETUNED_MODEL
+
+@pytest.fixture
+def pricer():
+    return Pricer()
+
+def test_wake_up():
+    pricer = Pricer()
+    assert pricer.wake_up() == "ok"
+
+
+@patch('transformers.AutoTokenizer')
+@patch('peft.PeftModel')
+@patch('transformers.AutoModelForCausalLM')
+def test_setup(MockAutoModel, MockPeftModel, MockAutoTokenizer, pricer):
+    # Setup mocks
+    mock_tokenizer = MockAutoTokenizer.from_pretrained.return_value
+    mock_model = MockAutoModel.from_pretrained.return_value
+    mock_peft_model = MockPeftModel.from_pretrained.return_value
+    
+    # Call the setup method
+    pricer.setup()
+    
+    # Assertions to ensure the setup method works correctly
+    MockAutoTokenizer.from_pretrained.assert_called_once_with(BASE_DIR)
+    assert pricer.tokenizer == mock_tokenizer
+    assert pricer.tokenizer.pad_token == pricer.tokenizer.eos_token
+    assert pricer.tokenizer.padding_side == "right"
+    
+    quant_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_quant_type="nf4"
+        )
+    
+    MockAutoModel.from_pretrained.assert_called_once_with(
+        BASE_DIR, 
+        quantization_config=quant_config, 
+        device_map="auto"
+        )
+    assert pricer.base_model == mock_model
+    
+    MockPeftModel.from_pretrained.assert_called_once_with(mock_model, FINETUNED_DIR, revision=REVISION)
+    assert pricer.fine_tuned_model == mock_peft_model
+
+
+@patch('transformers.AutoTokenizer')
+@patch('peft.PeftModel')
+def test_price(MockPeftModel, MockAutoTokenizer, pricer):
+ # Setup mocks
+    mock_tokenizer = MockAutoTokenizer.return_value
+    mock_tokenizer.encode.return_value = torch.tensor([[1, 2, 3]])
+    mock_tokenizer.decode.return_value = "Price is $123.45"
+    
+    mock_model = MockPeftModel.return_value
+    mock_model.generate.return_value = torch.tensor([[1, 2, 3, 4, 5]])
+    
+    # Assign mocks to the pricer instance
+    pricer.tokenizer = mock_tokenizer
+    pricer.fine_tuned_model = mock_model
+    
+    # Call the method
+    description = "Test description"
+    result = pricer.price(description)
+    
+    # Assert the result
+    assert result == 123.45