lfoppiano commited on
Commit
7bc374b
β€’
1 Parent(s): d74cacd

update to mistral v0.2, add selectable embeddings

Browse files
Files changed (1) hide show
  1. streamlit_app.py +52 -12
streamlit_app.py CHANGED
@@ -24,9 +24,23 @@ OPENAI_MODELS = ['gpt-3.5-turbo',
24
  "gpt-4",
25
  "gpt-4-1106-preview"]
26
 
 
 
 
 
 
 
27
  OPEN_MODELS = {
28
- 'mistral-7b-instruct-v0.1': 'mistralai/Mistral-7B-Instruct-v0.1',
29
  "zephyr-7b-beta": 'HuggingFaceH4/zephyr-7b-beta'
 
 
 
 
 
 
 
 
30
  }
31
 
32
  DISABLE_MEMORY = ['zephyr-7b-beta']
@@ -83,6 +97,9 @@ if 'pdf' not in st.session_state:
83
  if 'pdf_rendering' not in st.session_state:
84
  st.session_state['pdf_rendering'] = None
85
 
 
 
 
86
  st.set_page_config(
87
  page_title="Scientific Document Insights Q/A",
88
  page_icon="πŸ“",
@@ -139,24 +156,34 @@ def clear_memory():
139
 
140
 
141
  # @st.cache_resource
142
- def init_qa(model, api_key=None):
143
  ## For debug add: callbacks=[PromptLayerCallbackHandler(pl_tags=["langchain", "chatgpt", "document-qa"])])
144
  if model in OPENAI_MODELS:
 
 
 
145
  st.session_state['memory'] = ConversationBufferWindowMemory(k=4)
146
  if api_key:
147
  chat = ChatOpenAI(model_name=model,
148
  temperature=0,
149
  openai_api_key=api_key,
150
  frequency_penalty=0.1)
151
- embeddings = OpenAIEmbeddings(openai_api_key=api_key)
 
 
 
 
152
 
153
  else:
154
  chat = ChatOpenAI(model_name=model,
155
  temperature=0,
156
  frequency_penalty=0.1)
157
- embeddings = OpenAIEmbeddings()
158
 
159
  elif model in OPEN_MODELS:
 
 
 
160
  chat = HuggingFaceEndpoint(
161
  repo_id=OPEN_MODELS[model],
162
  temperature=0.01,
@@ -164,7 +191,7 @@ def init_qa(model, api_key=None):
164
  model_kwargs={"max_length": 4096}
165
  )
166
  embeddings = HuggingFaceEmbeddings(
167
- model_name="all-MiniLM-L6-v2")
168
  st.session_state['memory'] = ConversationBufferWindowMemory(k=4) if model not in DISABLE_MEMORY else None
169
  else:
170
  st.error("The model was not loaded properly. Try reloading. ")
@@ -231,15 +258,25 @@ with st.sidebar:
231
  "Model:",
232
  options=OPENAI_MODELS + list(OPEN_MODELS.keys()),
233
  index=(OPENAI_MODELS + list(OPEN_MODELS.keys())).index(
234
- "zephyr-7b-beta") if "DEFAULT_MODEL" not in os.environ or not os.environ["DEFAULT_MODEL"] else (
235
  OPENAI_MODELS + list(OPEN_MODELS.keys())).index(os.environ["DEFAULT_MODEL"]),
236
  placeholder="Select model",
237
  help="Select the LLM model:",
238
  disabled=st.session_state['doc_id'] is not None or st.session_state['uploaded']
239
  )
 
 
 
 
 
 
 
 
 
 
240
 
241
  st.markdown(
242
- ":warning: [Usage disclaimer](https://github.com/lfoppiano/document-qa/tree/review-interface#disclaimer-on-data-security-and-privacy-%EF%B8%8F) :warning: ")
243
 
244
  if (model in OPEN_MODELS) and model not in st.session_state['api_keys']:
245
  if 'HUGGINGFACEHUB_API_TOKEN' not in os.environ:
@@ -256,7 +293,7 @@ with st.sidebar:
256
  st.session_state['api_keys'][model] = api_key
257
  # if 'HUGGINGFACEHUB_API_TOKEN' not in os.environ:
258
  # os.environ["HUGGINGFACEHUB_API_TOKEN"] = api_key
259
- st.session_state['rqa'][model] = init_qa(model)
260
 
261
  elif model in OPENAI_MODELS and model not in st.session_state['api_keys']:
262
  if 'OPENAI_API_KEY' not in os.environ:
@@ -270,9 +307,9 @@ with st.sidebar:
270
  with st.spinner("Preparing environment"):
271
  st.session_state['api_keys'][model] = api_key
272
  if 'OPENAI_API_KEY' not in os.environ:
273
- st.session_state['rqa'][model] = init_qa(model, api_key)
274
  else:
275
- st.session_state['rqa'][model] = init_qa(model)
276
  # else:
277
  # is_api_key_provided = st.session_state['api_key']
278
 
@@ -371,10 +408,13 @@ with st.sidebar:
371
 
372
  st.header("Query mode (Advanced use)")
373
  st.markdown(
374
- """By default, the mode is set to LLM (Language Model) which enables question/answering. You can directly ask questions related to the document content, and the system will answer the question using content from the document.""")
 
375
 
376
  st.markdown(
377
- """If you switch the mode to "Embedding," the system will return specific chunks from the document that are semantically related to your query. This mode helps to test why sometimes the answers are not satisfying or incomplete. """)
 
 
378
 
379
  if uploaded_file and not st.session_state.loaded_embeddings:
380
  if model not in st.session_state['api_keys']:
 
24
  "gpt-4",
25
  "gpt-4-1106-preview"]
26
 
27
+ OPENAI_EMBEDDINGS = [
28
+ 'text-embedding-ada-002',
29
+ 'text-embedding-3-large',
30
+ 'openai-text-embedding-3-small'
31
+ ]
32
+
33
  OPEN_MODELS = {
34
+ 'mistral-7b-instruct-v0.2': 'mistralai/Mistral-7B-Instruct-v0.2',
35
  "zephyr-7b-beta": 'HuggingFaceH4/zephyr-7b-beta'
36
+ # 'Phi-3-mini-128k-instruct': "microsoft/Phi-3-mini-128k-instruct",
37
+ # 'Phi-3-mini-4k-instruct': "microsoft/Phi-3-mini-4k-instruct"
38
+ }
39
+
40
+ DEFAULT_OPEN_EMBEDDING_NAME = 'Default (all-MiniLM-L6-v2)'
41
+ OPEN_EMBEDDINGS = {
42
+ DEFAULT_OPEN_EMBEDDING_NAME: 'all-MiniLM-L6-v2',
43
+ 'Salesforce/SFR-Embedding-Mistral': 'Salesforce/SFR-Embedding-Mistral'
44
  }
45
 
46
  DISABLE_MEMORY = ['zephyr-7b-beta']
 
97
  if 'pdf_rendering' not in st.session_state:
98
  st.session_state['pdf_rendering'] = None
99
 
100
+ if 'embeddings' not in st.session_state:
101
+ st.session_state['embeddings'] = None
102
+
103
  st.set_page_config(
104
  page_title="Scientific Document Insights Q/A",
105
  page_icon="πŸ“",
 
156
 
157
 
158
  # @st.cache_resource
159
+ def init_qa(model, embeddings_name=None, api_key=None):
160
  ## For debug add: callbacks=[PromptLayerCallbackHandler(pl_tags=["langchain", "chatgpt", "document-qa"])])
161
  if model in OPENAI_MODELS:
162
+ if embeddings_name is None:
163
+ embeddings_name = 'text-embedding-ada-002'
164
+
165
  st.session_state['memory'] = ConversationBufferWindowMemory(k=4)
166
  if api_key:
167
  chat = ChatOpenAI(model_name=model,
168
  temperature=0,
169
  openai_api_key=api_key,
170
  frequency_penalty=0.1)
171
+ if embeddings_name not in OPENAI_EMBEDDINGS:
172
+ st.error(f"The embeddings provided {embeddings_name} are not supported by this model {model}.")
173
+ st.stop()
174
+ return
175
+ embeddings = OpenAIEmbeddings(model=embeddings_name, openai_api_key=api_key)
176
 
177
  else:
178
  chat = ChatOpenAI(model_name=model,
179
  temperature=0,
180
  frequency_penalty=0.1)
181
+ embeddings = OpenAIEmbeddings(model=embeddings_name)
182
 
183
  elif model in OPEN_MODELS:
184
+ if embeddings_name is None:
185
+ embeddings_name = DEFAULT_OPEN_EMBEDDING_NAME
186
+
187
  chat = HuggingFaceEndpoint(
188
  repo_id=OPEN_MODELS[model],
189
  temperature=0.01,
 
191
  model_kwargs={"max_length": 4096}
192
  )
193
  embeddings = HuggingFaceEmbeddings(
194
+ model_name=OPEN_EMBEDDINGS[embeddings_name])
195
  st.session_state['memory'] = ConversationBufferWindowMemory(k=4) if model not in DISABLE_MEMORY else None
196
  else:
197
  st.error("The model was not loaded properly. Try reloading. ")
 
258
  "Model:",
259
  options=OPENAI_MODELS + list(OPEN_MODELS.keys()),
260
  index=(OPENAI_MODELS + list(OPEN_MODELS.keys())).index(
261
+ "mistral-7b-instruct-v0.2") if "DEFAULT_MODEL" not in os.environ or not os.environ["DEFAULT_MODEL"] else (
262
  OPENAI_MODELS + list(OPEN_MODELS.keys())).index(os.environ["DEFAULT_MODEL"]),
263
  placeholder="Select model",
264
  help="Select the LLM model:",
265
  disabled=st.session_state['doc_id'] is not None or st.session_state['uploaded']
266
  )
267
+ embedding_choices = OPENAI_EMBEDDINGS if model in OPENAI_MODELS else OPEN_EMBEDDINGS
268
+
269
+ st.session_state['embeddings'] = embedding_name = st.selectbox(
270
+ "Embeddings:",
271
+ options=embedding_choices,
272
+ index=0,
273
+ placeholder="Select embedding",
274
+ help="Select the Embedding function:",
275
+ disabled=st.session_state['doc_id'] is not None or st.session_state['uploaded']
276
+ )
277
 
278
  st.markdown(
279
+ ":warning: [Usage disclaimer](https://github.com/lfoppiano/document-qa?tab=readme-ov-file#disclaimer-on-data-security-and-privacy-%EF%B8%8F) :warning: ")
280
 
281
  if (model in OPEN_MODELS) and model not in st.session_state['api_keys']:
282
  if 'HUGGINGFACEHUB_API_TOKEN' not in os.environ:
 
293
  st.session_state['api_keys'][model] = api_key
294
  # if 'HUGGINGFACEHUB_API_TOKEN' not in os.environ:
295
  # os.environ["HUGGINGFACEHUB_API_TOKEN"] = api_key
296
+ st.session_state['rqa'][model] = init_qa(model, embedding_name)
297
 
298
  elif model in OPENAI_MODELS and model not in st.session_state['api_keys']:
299
  if 'OPENAI_API_KEY' not in os.environ:
 
307
  with st.spinner("Preparing environment"):
308
  st.session_state['api_keys'][model] = api_key
309
  if 'OPENAI_API_KEY' not in os.environ:
310
+ st.session_state['rqa'][model] = init_qa(model, st.session_state['embeddings'], api_key)
311
  else:
312
+ st.session_state['rqa'][model] = init_qa(model, st.session_state['embeddings'])
313
  # else:
314
  # is_api_key_provided = st.session_state['api_key']
315
 
 
408
 
409
  st.header("Query mode (Advanced use)")
410
  st.markdown(
411
+ """By default, the mode is set to LLM (Language Model) which enables question/answering.
412
+ You can directly ask questions related to the document content, and the system will answer the question using content from the document.""")
413
 
414
  st.markdown(
415
+ """If you switch the mode to "Embedding," the system will return specific chunks from the document
416
+ that are semantically related to your query. This mode helps to test why sometimes the answers are not
417
+ satisfying or incomplete. """)
418
 
419
  if uploaded_file and not st.session_state.loaded_embeddings:
420
  if model not in st.session_state['api_keys']: