Daniel Marques commited on
Commit
4ec7545
1 Parent(s): 66a4e8f

feat: add stream

Browse files
Files changed (3) hide show
  1. main.py +7 -7
  2. run_localGPT.py +2 -2
  3. static/index.html +5 -1
main.py CHANGED
@@ -45,7 +45,7 @@ DB = Chroma(
45
 
46
  RETRIEVER = DB.as_retriever()
47
 
48
- LLM = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=False)
49
 
50
  template = """you are a helpful, respectful and honest assistant.
51
  Your name is Katara llma. You should only use the source documents provided to answer the questions.
@@ -134,12 +134,6 @@ def run_ingest_route():
134
  )
135
 
136
 
137
- generated_text = ""
138
- for new_text in StreamData:
139
- generated_text += new_text
140
- print(generated_text)
141
-
142
-
143
  return {"response": "The training was successfully completed"}
144
  except Exception as e:
145
  raise HTTPException(status_code=500, detail=f"Error occurred: {str(e)}")
@@ -189,6 +183,12 @@ async def predict(data: Predict):
189
  (os.path.basename(str(document.metadata["source"])), str(document.page_content))
190
  )
191
 
 
 
 
 
 
 
192
  return {"response": prompt_response_dict}
193
  else:
194
  raise HTTPException(status_code=400, detail="Prompt Incorrect")
 
45
 
46
  RETRIEVER = DB.as_retriever()
47
 
48
+ LLM, STREAMER = load_model(device_type=DEVICE_TYPE, model_id=MODEL_ID, model_basename=MODEL_BASENAME, stream=False)
49
 
50
  template = """you are a helpful, respectful and honest assistant.
51
  Your name is Katara llma. You should only use the source documents provided to answer the questions.
 
134
  )
135
 
136
 
 
 
 
 
 
 
137
  return {"response": "The training was successfully completed"}
138
  except Exception as e:
139
  raise HTTPException(status_code=500, detail=f"Error occurred: {str(e)}")
 
183
  (os.path.basename(str(document.metadata["source"])), str(document.page_content))
184
  )
185
 
186
+
187
+ generated_text = ""
188
+ for new_text in STREAMER:
189
+ generated_text += new_text
190
+ print(generated_text)
191
+
192
  return {"response": prompt_response_dict}
193
  else:
194
  raise HTTPException(status_code=400, detail="Prompt Incorrect")
run_localGPT.py CHANGED
@@ -89,13 +89,13 @@ def load_model(device_type, model_id, model_basename=None, LOGGING=logging, stre
89
  top_k=40,
90
  repetition_penalty=1.0,
91
  generation_config=generation_config,
92
- callback=[StreamingStdOutCallbackHandler()]
93
  )
94
 
95
  local_llm = HuggingFacePipeline(pipeline=pipe)
96
  logging.info("Local LLM Loaded")
97
 
98
- return local_llm
99
 
100
 
101
  def retrieval_qa_pipline(device_type, use_history, promptTemplate_type="llama"):
 
89
  top_k=40,
90
  repetition_penalty=1.0,
91
  generation_config=generation_config,
92
+ streamer=streamer
93
  )
94
 
95
  local_llm = HuggingFacePipeline(pipeline=pipe)
96
  logging.info("Local LLM Loaded")
97
 
98
+ return (local_llm, streamer)
99
 
100
 
101
  def retrieval_qa_pipline(device_type, use_history, promptTemplate_type="llama"):
static/index.html CHANGED
@@ -6,7 +6,11 @@
6
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
7
  <!-- Bootstrap CSS -->
8
 
9
- <!-- <link rel="icon" type="image" href="static\social_icons\favicon.png" /> -->
 
 
 
 
10
  <script src="https://code.jquery.com/jquery-3.6.0.min.js"></script>
11
  <link
12
  href="https://cdn.jsdelivr.net/npm/bootstrap@5.1.3/dist/css/bootstrap.min.css"
 
6
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
7
  <!-- Bootstrap CSS -->
8
 
9
+ <link
10
+ rel="icon"
11
+ type="image"
12
+ href="https://katara.earth/katara-favicon.svg"
13
+ />
14
  <script src="https://code.jquery.com/jquery-3.6.0.min.js"></script>
15
  <link
16
  href="https://cdn.jsdelivr.net/npm/bootstrap@5.1.3/dist/css/bootstrap.min.css"