fix: add streamer

load_models.py

@@ -222,9 +222,4 @@ def load_model(device_type, model_id, model_basename=None, LOGGING=logging, stre
     local_llm = HuggingFacePipeline(pipeline=pipe)
     logging.info("Local LLM Loaded")
 
-    generated_text = ""
-    for new_text in streamer:
-        generated_text += new_text
-        print(generated_text)
-
-    return local_llm
+    return [local_llm, streamer]

main.py

@@ -42,9 +42,11 @@ DB = Chroma(
 
 RETRIEVER = DB.as_retriever()
 
-LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=False)
+models = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=False)
+LLM = models[0]
+STREAMER = models[1]
 
-template = """Your name is Katara and you are a helpful, respectful and honest assistant. You should only use the source documents provided to answer the questions.
+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.
@@ -70,7 +72,6 @@ QA = RetrievalQA.from_chain_type(
     },
 )
 
-
 class Predict(BaseModel):
     prompt: str
 
@@ -145,7 +146,7 @@ def get_files():
 def delete_source_route(data: Delete):
     filename = data.filename
     path_source_documents = os.path.join(os.getcwd(), PATH_NAME_SOURCE_DIRECTORY)
-    file_to_delete = f"{path_source_documents}/${filename}"
+    file_to_delete = f"{path_source_documents}/{filename}"
 
     if os.path.exists(file_to_delete):
         try:
@@ -166,6 +167,9 @@ async def predict(data: Predict):
         # print(f'User Prompt: {user_prompt}')
         # Get the answer from the chain
         res = QA(user_prompt)
+
+        print(res)
+
         answer, docs = res["result"], res["source_documents"]
 
         prompt_response_dict = {
@@ -179,6 +183,11 @@ 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)
+
         return {"response": prompt_response_dict}
     else:
         raise HTTPException(status_code=400, detail="Prompt Incorrect")