jalr/tabbyAPI-ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Tests: Fix linting

Browse source 
Also change how wheel_test works for safe import testing rather than
trying to import the package which can cause system issues.

Signed-off-by: kingbri <bdashore3@proton.me>
This commit is contained in:
kingbri 2023-12-13 23:05:50 -05:00
parent 036ba2669c
commit 59729e2a4a
2 changed files with 18 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

3
tests/model_test.py
View file

@ -9,7 +9,8 @@ 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 = "")
for g in generator:
print(g, end = "")
container.unload()
del container

38
tests/wheel_test.py
View file

@ -1,47 +1,41 @@
import traceback
from importlib.metadata import version
from importlib.util import find_spec
successful_packages = []
errored_packages = []
try:
import flash_attn
if find_spec("flash_attn") is not None:
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())
else:
print("Flash attention 2 is not found in your environment.")
errored_packages.append("flash_attn")
try:
import exllamav2
if find_spec("exllamav2") is not None:
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())
else:
print("Exllamav2 is not found in your environment.")
errored_packages.append("exllamav2")
try:
import torch
if find_spec("torch") is not None:
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())
else:
print("Torch is not found in your environment.")
errored_packages.append("torch")
try:
import fastchat
if find_spec("fastchat") is not None:
print(f"Fastchat on version {version('fschat')} 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())
else:
print("Fastchat is not found in your environment. It isn't needed unless you're using chat completions with message arrays.")
errored_packages.append("fastchat")
print(
f"\nSuccessful imports: {', '.join(successful_packages)}",
f"\nErrored imports: {''.join(errored_packages)}"
)
if len(errored_packages) > 0:
print("\nIf packages are installed, but not found on this test, please check the wheel versions for the correct python version and CUDA version (if applicable).")