added tests

requirements.txt

@@ -7,4 +7,5 @@ beautifulsoup4~=4.11.2
 urllib3
 requests~=2.28.2
 accelerate
-bitsandbytes
+bitsandbytes
+pytest

tests.py

@@ -0,0 +1,43 @@
+import os
+import shutil
+
+import pytest as pytest
+from grouped_sampling import GroupedSamplingPipeLine
+
+from hanlde_form_submit import create_pipeline
+from prompt_engeneering import rewrite_prompt
+from supported_models import get_supported_model_names
+
+HUGGING_FACE_CACHE_DIR = "/home/yoni/.cache/huggingface/hub"
+
+
+def test_prompt_engineering():
+    example_prompt = "Answer yes or no, is the sky blue?"
+    rewritten_prompt = rewrite_prompt(example_prompt)
+    assert rewritten_prompt.startswith("Web search results:")
+    assert rewritten_prompt.endswith("Query: Answer yes or no, is the sky blue?")
+    assert "Current date: " in rewritten_prompt
+    assert "Instructions: " in rewritten_prompt
+
+
+def test_get_supported_model_names():
+    supported_model_names = get_supported_model_names()
+    assert len(supported_model_names) > 0
+    assert "gpt2" in supported_model_names
+    assert all([isinstance(name, str) for name in supported_model_names])
+
+
+@pytest.mark.parametrize("model_name", get_supported_model_names())
+def test_create_pipeline(model_name: str):
+    pipeline: GroupedSamplingPipeLine = create_pipeline(model_name, 5)
+    assert pipeline is not None
+    assert pipeline.model_name == model_name
+    assert pipeline.wrapped_model.group_size == 5
+    assert pipeline.wrapped_model.end_of_sentence_stop is True
+    model_folder_name = "models--" + model_name.replace("/", "--")
+    model_path = os.path.join(HUGGING_FACE_CACHE_DIR, model_folder_name)
+    shutil.rmtree(model_path)
+
+
+if __name__ == "__main__":
+    pytest.main()