fix: add callback StreamingStdOutCallbackHandler

load_models.py

@@ -5,6 +5,7 @@ import logging
 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 transformers import (
@@ -68,7 +69,7 @@ def load_quantized_model_gguf_ggml(model_id, model_basename, device_type, loggin
         kwargs["stream"] = stream
 
         if stream == True:
-            kwargs["callbacks"] = callbacks
+            kwargs["callbacks"] = [StreamingStdOutCallbackHandler()]
 
         return LlamaCpp(**kwargs)
     except:
@@ -220,7 +221,6 @@ def load_model(device_type, model_id, model_basename=None, LOGGING=logging, stre
         top_k=40,
         repetition_penalty=1.0,
         generation_config=generation_config,
-        # callbacks=callbacks
     )
 
     local_llm = HuggingFacePipeline(pipeline=pipe)

main.py

@@ -44,19 +44,6 @@ class MyCustomHandler(BaseCallbackHandler):
     async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         print("finish")
 
-class CustomHandler(BaseCallbackHandler):
-    def on_llm_new_token(self, token: str, **kwargs) -> None:
-        print(f" CustomHandler: {token}")
-
-    async def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> None:
-        class_name = serialized["name"]
-        print("CustomHandler start")
-
-    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        print("CustomHandler finish")
-
 # if torch.backends.mps.is_available():
 #     DEVICE_TYPE = "mps"
 # elif torch.cuda.is_available():
@@ -101,7 +88,6 @@ QA = RetrievalQA.from_chain_type(
     chain_type_kwargs={
         "prompt": QA_CHAIN_PROMPT,
         "memory": memory,
-        "callbacks": [CustomHandler()]
     },
 )