jharrison27 commited on
Commit
bae5df9
1 Parent(s): 3dbe475

Handle no colbert indexes

Browse files
Files changed (1) hide show
  1. app.py +26 -14
app.py CHANGED
@@ -20,15 +20,22 @@ GENERATE_KWARGS = {
20
  "do_sample": False,
21
  }
22
 
23
- # RAG Model setup
24
- RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")
25
-
26
  try:
27
- gr.Info("Setting up retriever, please wait...")
28
- rag_initial_output = RAG.search("What is Generative AI in Healthcare?", k=1)
29
- gr.Info("Retriever working successfully!")
30
- except Exception as e:
31
- gr.Warning(f"Retriever not working: {str(e)}")
 
 
 
 
 
 
 
 
 
 
32
 
33
  # Header setup
34
  mark_text = '# 🩺🔍 Search Results\n'
@@ -46,7 +53,10 @@ try:
46
  except FileNotFoundError:
47
  index_info = "Semantic Search"
48
 
49
- database_choices = [index_info, 'Arxiv Search - Latest - (EXPERIMENTAL)']
 
 
 
50
 
51
  # Arxiv API setup
52
  arx_client = arxiv.Client()
@@ -77,10 +87,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
77
 
78
  def update_with_rag_md(search_query, llm_results_use=5, database_choice=index_info, llm_model_picked=DEFAULT_LLM_MODEL):
79
  prompt_text_from_data = ""
80
- database_to_use = database_choice
81
 
82
- if database_choice == index_info:
83
  rag_out = get_rag(search_query, RAG, RETRIEVE_RESULTS)
 
84
  else:
85
  arxiv_search_success = True
86
  try:
@@ -89,11 +99,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
89
  arxiv_search_success = False
90
  except Exception as e:
91
  arxiv_search_success = False
92
- gr.Warning(f"Arxiv Search not working: {str(e)}, switching to semantic search ...")
93
 
94
  if not arxiv_search_success:
95
- rag_out = get_rag(search_query, RAG, RETRIEVE_RESULTS)
96
- database_to_use = index_info
 
 
97
 
98
  md_text_updated = mark_text
99
  for i, rag_answer in enumerate(rag_out):
 
20
  "do_sample": False,
21
  }
22
 
 
 
 
23
  try:
24
+ # RAG Model setup
25
+ RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")
26
+ semantic_search_available = True
27
+
28
+ try:
29
+ gr.Info("Setting up retriever, please wait...")
30
+ rag_initial_output = RAG.search("What is Generative AI in Healthcare?", k=1)
31
+ gr.Info("Retriever working successfully!")
32
+ except Exception as e:
33
+ gr.Warning(f"Retriever not working: {str(e)}")
34
+
35
+ except FileNotFoundError:
36
+ RAG = None
37
+ semantic_search_available = False
38
+ gr.Warning("Colbert index not found. Semantic search will be unavailable.")
39
 
40
  # Header setup
41
  mark_text = '# 🩺🔍 Search Results\n'
 
53
  except FileNotFoundError:
54
  index_info = "Semantic Search"
55
 
56
+ if semantic_search_available:
57
+ database_choices = [index_info, 'Arxiv Search - Latest']
58
+ else:
59
+ database_choices = ['Arxiv Search - Latest']
60
 
61
  # Arxiv API setup
62
  arx_client = arxiv.Client()
 
87
 
88
  def update_with_rag_md(search_query, llm_results_use=5, database_choice=index_info, llm_model_picked=DEFAULT_LLM_MODEL):
89
  prompt_text_from_data = ""
 
90
 
91
+ if database_choice == index_info and semantic_search_available:
92
  rag_out = get_rag(search_query, RAG, RETRIEVE_RESULTS)
93
+ database_to_use = 'Semantic Search'
94
  else:
95
  arxiv_search_success = True
96
  try:
 
99
  arxiv_search_success = False
100
  except Exception as e:
101
  arxiv_search_success = False
102
+ gr.Warning(f"Arxiv Search not working: {str(e)}")
103
 
104
  if not arxiv_search_success:
105
+ gr.Warning("Arxiv search failed. Please try again later.")
106
+ return "", ""
107
+
108
+ database_to_use = 'Arxiv Search'
109
 
110
  md_text_updated = mark_text
111
  for i, rag_answer in enumerate(rag_out):