50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
|
|
""" Test that GGUF models are loading correctly in local environment. By default, will run through a series of
|
|
different GGUF models in the ModelCatalog to spot-check that the model is correctly loading and
|
|
successfully completing an inference:
|
|
|
|
# tests several different underlying models:
|
|
|
|
# bling-answer-tool -> tiny-llama (1b)
|
|
# bling-phi-3-gguf -> phi-3 (3.8b)
|
|
# dragon-yi-answer-tool -> yi (6b)
|
|
# dragon-llama-answer-tool -> llama-2 (7b)
|
|
# llama-2-7b-chat-gguf -> llama-2-chat (7b)
|
|
# dragon-mistral-answer-tool -> mistral-1 (7b)
|
|
|
|
"""
|
|
|
|
|
|
from llmware.models import ModelCatalog
|
|
|
|
|
|
def test_gguf_model_load():
|
|
|
|
# feel free to adapt this model list
|
|
|
|
model_list = ["bling-answer-tool",
|
|
"bling-phi-3-gguf",
|
|
"dragon-yi-answer-tool",
|
|
"dragon-llama-answer-tool",
|
|
"llama-2-7b-chat-gguf",
|
|
"dragon-mistral-answer-tool"]
|
|
|
|
# please note that the unusually short and simple prompt at times actually yields more variability in the model
|
|
# response - we are only testing for successful loading and inference
|
|
|
|
sample_prompt = ("The company stock declined by $12 after poor earnings results."
|
|
"\nHow much did the stock price decline?")
|
|
|
|
for model_name in model_list:
|
|
|
|
print("\nmodel name: ", model_name)
|
|
|
|
model = ModelCatalog().load_model(model_name, temperature=0.0, sample=False)
|
|
|
|
response = model.inference(sample_prompt)
|
|
|
|
print(f"{model_name} - response: ", response)
|
|
|
|
assert response is not None
|
|
|