fix: add streamer

load_models.py

@@ -1,9 +1,15 @@
 import torch
+import asyncio
 import logging
+from typing import Any, Dict, List
+
 from auto_gptq import AutoGPTQForCausalLM
 from huggingface_hub import hf_hub_download
 from langchain.llms import LlamaCpp, HuggingFacePipeline
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.schema import LLMResult
+
+from langchain.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler
 
 from transformers import (
     AutoModelForCausalLM,
@@ -22,6 +28,29 @@ torch.set_grad_enabled(False)
 from constants import CONTEXT_WINDOW_SIZE, MAX_NEW_TOKENS, N_GPU_LAYERS, N_BATCH, MODELS_PATH
 
 
+class MyCustomSyncHandler(BaseCallbackHandler):
+    def on_llm_new_token(self, token: str, **kwargs) -> None:
+        print(f"Sync handler being called in a `thread_pool_executor`: token: {token}")
+
+class MyCustomAsyncHandler(AsyncCallbackHandler):
+    """Async callback handler that can be used to handle callbacks from langchain."""
+
+    async def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> None:
+        """Run when chain starts running."""
+        print("zzzz....")
+        await asyncio.sleep(0.3)
+        class_name = serialized["name"]
+        print("Hi! I just woke up. Your llm is starting")
+
+    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        """Run when chain ends running."""
+        print("zzzz....")
+        await asyncio.sleep(0.3)
+        print("Hi! I just woke up. Your llm is ending")
+
+
 def load_quantized_model_gguf_ggml(model_id, model_basename, device_type, logging, stream = False):
     """
     Load a GGUF/GGML quantized model using LlamaCpp.
@@ -66,6 +95,7 @@ def load_quantized_model_gguf_ggml(model_id, model_basename, device_type, loggin
 
         #add stream
         kwargs["stream"] = stream
+        kwargs["callbacks"] = [MyCustomSyncHandler(), MyCustomAsyncHandler()]
 
         return LlamaCpp(**kwargs)
     except:
@@ -220,7 +250,7 @@ def load_model(device_type, model_id, model_basename=None, LOGGING=logging, stre
         repetition_penalty=1.0,
         generation_config=generation_config,
         streamer=streamer,
-        callbacks=[StreamingStdOutCallbackHandler()]
+        callbacks=[MyCustomSyncHandler(), MyCustomAsyncHandler()]
     )
 
     local_llm = HuggingFacePipeline(pipeline=pipe)

main.py

@@ -42,10 +42,7 @@ DB = Chroma(
 
 RETRIEVER = DB.as_retriever()
 
-models = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True)
-
-LLM = models[0]
-STREAMER = models[1]
+LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True)
 
 template = """you are a helpful, respectful and honest assistant. You should only use the source documents provided to answer the questions.
 You should only respond only topics that contains in documents use to training.
@@ -182,10 +179,10 @@ async def predict(data: Predict):
                 (os.path.basename(str(document.metadata["source"])), str(document.page_content))
             )
 
-        generated_text = ""
-        for new_text in STREAMER:
-            generated_text += new_text
-            print(generated_text)
+        # generated_text = ""
+        # for new_text in STREAMER:
+        #     generated_text += new_text
+        #     print(generated_text)
 
         return {"response": prompt_response_dict}
     else: