Tests: Create

Add a few tests for the user to check if stuff works.

Signed-off-by: kingbri <bdashore3@proton.me>

tests/model_test.py

@@ -0,0 +1,21 @@
+from model import ModelContainer
+
+def progress(module, modules):
+    yield module, modules
+
+container = ModelContainer("/mnt/str/models/_exl2/mistral-7b-instruct-exl2/4.0bpw/")
+loader = container.load_gen(progress)
+for (module, modules) in loader:
+    print(module, modules)
+
+generator = container.generate_gen("Once upon a tim", token_healing = True)
+for g in generator: print(g, end = "")
+
+container.unload()
+del container
+
+mc = ModelContainer("/mnt/str/models/_exl2/mistral-7b-instruct-exl2/4.65bpw/")
+mc.load(progress)
+
+response = mc.generate("All work and no play makes turbo a derpy cat.\nAll work and no play makes turbo a derpy cat.\nAll", top_k = 1, max_new_tokens = 1000, stream_interval = 0.5)
+print (response)

tests/wheel_test.py

@@ -0,0 +1,47 @@
+import traceback
+from importlib.metadata import version
+
+successful_packages = []
+errored_packages = []
+
+try:
+    import flash_attn
+    print(f"Flash attention on version {version('flash_attn')} successfully imported")
+    successful_packages.append("flash_attn")
+except:
+    print("Flash attention could not be loaded because:")
+    print(traceback.format_exc())
+    errored_packages.append("flash_attn")
+
+try:
+    import exllamav2
+    print(f"Exllamav2 on version {version('exllamav2')} successfully imported")
+    successful_packages.append("exllamav2")
+except:
+    print("Exllamav2 could not be loaded because:")
+    print(traceback.format_exc())
+    errored_packages.append("exllamav2")
+
+try:
+    import torch
+    print(f"Torch on version {version('torch')} successfully imported")
+    successful_packages.append("torch")
+except:
+    print("Torch could not be loaded because:")
+    print(traceback.format_exc())
+    errored_packages.append("torch")
+
+try:
+    import fastchat
+    print(f"Fastchat on version {version('fastchat')} successfully imported")
+    successful_packages.append("fastchat")
+except:
+    print("Fastchat is only needed for chat completions with message arrays. Ignore this error if this isn't your usecase.")
+    print("Fastchat could not be loaded because:")
+    print(traceback.format_exc())
+    errored_packages.append("fastchat")
+
+print(
+    f"\nSuccessful imports: {', '.join(successful_packages)}",
+    f"\nErrored imports: {''.join(errored_packages)}"
+)