fix: add streamer

load_models.py

@@ -1,5 +1,5 @@
 import torch
-import asyncio
+
 import logging
 from typing import Any, Dict, List
 
@@ -7,9 +7,6 @@ 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,
@@ -28,30 +25,7 @@ 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):
+def load_quantized_model_gguf_ggml(model_id, model_basename, device_type, logging, stream = False, callbacks = []):
     """
     Load a GGUF/GGML quantized model using LlamaCpp.
 
@@ -93,9 +67,10 @@ def load_quantized_model_gguf_ggml(model_id, model_basename, device_type, loggin
         if device_type.lower() == "cuda":
             kwargs["n_gpu_layers"] = N_GPU_LAYERS  # set this based on your GPU
 
-        #add stream
         kwargs["stream"] = stream
-        kwargs["callbacks"] = [MyCustomSyncHandler(), MyCustomAsyncHandler()]
+
+        if stream == True:
+            kwargs["callbacks"] = callbacks
 
         return LlamaCpp(**kwargs)
     except:
@@ -145,6 +120,7 @@ def load_quantized_model_qptq(model_id, model_basename, device_type, logging):
         use_triton=False,
         quantize_config=None,
     )
+
     return model, tokenizer
 
 
@@ -195,7 +171,7 @@ def load_full_model(model_id, model_basename, device_type, logging):
     return model, tokenizer
 
 
-def load_model(device_type, model_id, model_basename=None, LOGGING=logging, stream=False):
+def load_model(device_type, model_id, model_basename=None, LOGGING=logging, stream=False, callbacks = []):
     """
     Select a model for text generation using the HuggingFace library.
     If you are running this for the first time, it will download a model for you.
@@ -219,7 +195,7 @@ def load_model(device_type, model_id, model_basename=None, LOGGING=logging, stre
 
     if model_basename is not None:
         if ".gguf" in model_basename.lower():
-            llm = load_quantized_model_gguf_ggml(model_id, model_basename, device_type, LOGGING, stream)
+            llm = load_quantized_model_gguf_ggml(model_id, model_basename, device_type, LOGGING, stream, callbacks)
             return llm
         elif ".ggml" in model_basename.lower():
             model, tokenizer = load_quantized_model_gguf_ggml(model_id, model_basename, device_type, LOGGING)

main.py

@@ -1,17 +1,21 @@
-from fastapi import FastAPI, HTTPException, UploadFile, WebSocket
-from fastapi.staticfiles import StaticFiles
-
-from pydantic import BaseModel
 import os
 import glob
 import shutil
 import subprocess
+import asyncio
+
+from fastapi import FastAPI, HTTPException, UploadFile, WebSocket
+from fastapi.staticfiles import StaticFiles
+
+from pydantic import BaseModel
 
 # import torch
 from langchain.chains import RetrievalQA
 from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.prompts import PromptTemplate
 from langchain.memory import ConversationBufferMemory
+from langchain.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler
+from langchain.schema import LLMResult
 
 # from langchain.embeddings import HuggingFaceEmbeddings
 from load_models import load_model
@@ -21,6 +25,26 @@ from langchain.vectorstores import Chroma
 
 from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY, MODEL_ID, MODEL_BASENAME, PATH_NAME_SOURCE_DIRECTORY
 
+class Predict(BaseModel):
+    prompt: str
+
+class Delete(BaseModel):
+    filename: str
+
+
+class MyCustomAsyncHandler(AsyncCallbackHandler):
+    def on_llm_new_token(self, token: str, **kwargs) -> None:
+        print(f" token: {token}")
+
+    async def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> None:
+        class_name = serialized["name"]
+        print("start")
+
+    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        print("finish")
+
 # if torch.backends.mps.is_available():
 #     DEVICE_TYPE = "mps"
 # elif torch.cuda.is_available():
@@ -42,15 +66,13 @@ DB = Chroma(
 
 RETRIEVER = DB.as_retriever()
 
-LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True)
+LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=True, callbacks = [MyCustomAsyncHandler])
 
-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.
-Use the following pieces of context to answer the question at the end.
-Always answer in the most helpful and safe way possible.
-If you don't know the answer to a question, just say that you don't know, don't try to make up an answer, don't share false information.
-Use 15 sentences maximum. Keep the answer as concise as possible.
-Always say "thanks for asking!" at the end of the answer.
+template = """you are a helpful, respectful and honest assistant. When answering questions, you should only use the documents provided.
+You should only answer the topics that appear in these documents.
+Always answer in the most helpful and reliable way possible, if you don't know the answer to a question, just say you don't know, don't try to make up an answer,
+don't share false information. you should use no more than 15 sentences and all your answers should be as concise as possible.
+Always say "Thank you for asking!" at the end of your answer.
 Context: {history} \n {context}
 Question: {question}
 """
@@ -70,12 +92,6 @@ QA = RetrievalQA.from_chain_type(
     },
 )
 
-class Predict(BaseModel):
-    prompt: str
-
-class Delete(BaseModel):
-    filename: str
-
 app = FastAPI(title="homepage-app")
 api_app = FastAPI(title="api app")
 
@@ -179,6 +195,12 @@ async def predict(data: Predict):
                 (os.path.basename(str(document.metadata["source"])), str(document.page_content))
             )
 
+        qa_chain_response = res.stream(
+            {"query": user_prompt},
+        )
+
+        print(f"{qa_chain_response} stream")
+
         # generated_text = ""
         # for new_text in STREAMER:
         #     generated_text += new_text