rag-chatbot-2

1
import time
2
from pathlib import Path
3

4
from exp_lama_cpp.model import Model, get_model_setting
5

6
# CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install llama-cpp-python
7

8
if __name__ == "__main__":
9
    root_folder = Path(__file__).resolve().parent.parent.parent
10
    model_folder = root_folder / "models"
11
    Path(model_folder).parent.mkdir(parents=True, exist_ok=True)
12

13
    model_settings = get_model_setting("stablelm-zephyr")
14

15
    llm = Model(model_folder, model_settings)
16

17
    start_time = time.time()
18
    prompt = llm.generate_summarization_prompt(text="<put the text here>")
19
    output = llm.generate_answer(prompt, max_new_tokens=512)
20
    print(output)
21
    took = time.time() - start_time
22
    print(f"\n--- Took {took:.2f} seconds ---")
23

24
    start_time = time.time()
25
    stream = llm.start_answer_iterator_streamer(prompt, max_new_tokens=256)
26
    for output in stream:
27
        print(output["choices"][0]["text"], end="", flush=True)
28
    took = time.time() - start_time
29

30
    print(f"\n--- Took {took:.2f} seconds ---")
31