fix: add console trupple

.flake8

@@ -1,4 +0,0 @@
-[flake8]
-exclude = docs
-max-line-length = 119
-extend-ignore = E203

load_models.py

@@ -204,6 +204,8 @@ def load_model(device_type, model_id, model_basename=None, LOGGING=logging, stre
 
     streamer = TextStreamer(tokenizer, skip_prompt=True)
 
+    logging.info(streamer)
+
     pipe = pipeline(
         "text-generation",
         model=model,
@@ -220,4 +222,9 @@ def load_model(device_type, model_id, model_basename=None, LOGGING=logging, stre
     local_llm = HuggingFacePipeline(pipeline=pipe)
     logging.info("Local LLM Loaded")
 
-    return local_llm, streamer
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        print(generated_text)
+
+    return local_llm

main.py

@@ -1,7 +1,6 @@
 from fastapi import FastAPI, HTTPException, UploadFile, WebSocket
 from fastapi.staticfiles import StaticFiles
 
-
 from pydantic import BaseModel
 import os
 import glob
@@ -14,7 +13,6 @@ from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.prompts import PromptTemplate
 from langchain.memory import ConversationBufferMemory
 
-
 # from langchain.embeddings import HuggingFaceEmbeddings
 from load_models import load_model
 
@@ -44,11 +42,7 @@ DB = Chroma(
 
 RETRIEVER = DB.as_retriever()
 
-models = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=False)
-
-print(models)
-
-LLM, STREAMER = models
+LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=False)
 
 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.
 You should only respond only topics that contains in documents use to training.
@@ -186,10 +180,7 @@ async def predict(data: Predict):
             )
 
 
-        generated_text = ""
-        for new_text in STREAMER:
-            generated_text += new_text
-            print(generated_text)
+
 
         return {"response": prompt_response_dict}
     else:

requirements.txt

@@ -24,7 +24,6 @@ accelerate
 bitsandbytes ; sys_platform != 'win32'
 bitsandbytes-windows ; sys_platform == 'win32'
 click
-flask
 requests
 uvicorn
 fastapi