File size: 2,567 Bytes
7f9d16c
815128e
c2cb992
7f9d16c
c2cb992
815128e
7f9d16c
 
815128e
e182c41
7f9d16c
e182c41
815128e
c2cb992
e182c41
 
815128e
e182c41
 
815128e
e182c41
 
 
815128e
7f9d16c
 
c2cb992
 
 
7f9d16c
 
815128e
7f9d16c
 
 
 
 
 
815128e
c2cb992
e182c41
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
815128e
c2cb992
e182c41
815128e
c2cb992
e182c41
815128e
7f9d16c
e182c41
815128e
 
7f9d16c
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
# project/test.py

import os
import unittest
from timeit import default_timer as timer

from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import HumanMessage

from app_modules.init import app_init
from app_modules.llm_loader import LLMLoader
from app_modules.utils import get_device_types, print_llm_response


class TestLLMLoader:  # (unittest.TestCase):
    question = "What's the capital city of Malaysia?"

    def run_test_case(self, llm_model_type, query):
        n_threds = int(os.environ.get("NUMBER_OF_CPU_CORES") or "4")

        hf_embeddings_device_type, hf_pipeline_device_type = get_device_types()
        print(f"hf_embeddings_device_type: {hf_embeddings_device_type}")
        print(f"hf_pipeline_device_type: {hf_pipeline_device_type}")

        llm_loader = LLMLoader(llm_model_type)
        start = timer()
        llm_loader.init(
            n_threds=n_threds, hf_pipeline_device_type=hf_pipeline_device_type
        )
        end = timer()
        print(f"Model loaded in {end - start:.3f}s")

        result = llm_loader.llm(
            [HumanMessage(content=query)] if llm_model_type == "openai" else query
        )
        end2 = timer()
        print(f"Inference completed in {end2 - end:.3f}s")
        print(result)

    def test_openai(self):
        self.run_test_case("openai", self.question)

    def test_llamacpp(self):
        self.run_test_case("llamacpp", self.question)

    def test_gpt4all_j(self):
        self.run_test_case("gpt4all-j", self.question)

    def test_huggingface(self):
        self.run_test_case("huggingface", self.question)


class TestQAChain(unittest.TestCase):
    qa_chain: any
    question = "What's deep learning?"

    def run_test_case(self, llm_model_type, query):
        start = timer()
        os.environ["LLM_MODEL_TYPE"] = llm_model_type
        qa_chain = app_init()
        end = timer()
        print(f"App initialized in {end - start:.3f}s")

        inputs = {"question": query, "chat_history": []}
        result = qa_chain.call_chain(inputs, None)
        end2 = timer()
        print(f"Inference completed in {end2 - end:.3f}s")
        print_llm_response(result)

    def test_openai(self):
        self.run_test_case("openai", self.question)

    def test_llamacpp(self):
        self.run_test_case("llamacpp", self.question)

    def test_gpt4all_j(self):
        self.run_test_case("gpt4all-j", self.question)

    def test_huggingface(self):
        self.run_test_case("huggingface", self.question)


if __name__ == "__main__":
    unittest.main()