Added experimental Arxiv Support

README.md

@@ -11,4 +11,4 @@ license: cc0-1.0
 ---
 
 ## Arxiv-CS-RAG
-Index Last Updated : 2024-03-03
+Index Last Updated : 2024-03-10

app.py

@@ -7,6 +7,9 @@ import re
 from datetime import datetime
 import json
 
+import arxiv
+from utils import get_md_text_abstract, search_cleaner, get_arxiv_live_search
+
 retrieve_results = 10 
 show_examples = False
 llm_models_to_choose = ['mistralai/Mixtral-8x7B-Instruct-v0.1','mistralai/Mistral-7B-Instruct-v0.2', 'google/gemma-7b-it', 'None']
@@ -18,6 +21,7 @@ generate_kwargs = dict(
     do_sample = False,
     )
 
+## RAG Model
 RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")
 
 try:
@@ -28,8 +32,9 @@ try:
 except:
   gr.Warning("Retriever not working!")
 
+## Header
 mark_text = '# 🔍 Search Results\n'
-header_text = "# ArXivCS RAG \n"
+header_text = "# ArXiv CS RAG \n"
 
 try:
   with open("README.md", "r") as f:
@@ -39,10 +44,24 @@ try:
   date = match.group().split(': ')[1]
   formatted_date = datetime.strptime(date, '%Y-%m-%d').strftime('%d %b %Y')
   header_text += f'Index Last Updated: {formatted_date}\n'
-    
+  index_info = f"Semantic Search - up to {formatted_date}"  
 except:
-  pass
+  index_info = "Semantic Search"
+
+database_choices = [index_info,'Arxiv Search - Latest - (EXPERIMENTAL)']
+
+## Arxiv API
+arx_client = arxiv.Client()
+is_arxiv_available = True
+check_arxiv_result = get_arxiv_live_search("What is Mistral?", arx_client, retrieve_results)
+if len(check_arxiv_result) == 0:
+  is_arxiv_available = False
+  print("Arxiv search not working, switching to default search ...")
+  database_choices = [index_info]
+
 
+
+## Show examples (disabled)
 if show_examples:
     with open("sample_outputs.json", "r") as f:
       sample_outputs = json.load(f)
@@ -91,27 +110,43 @@ with gr.Blocks(theme = gr.themes.Soft()) as demo:
         with gr.Row(equal_height = True):
           llm_model = gr.Dropdown(choices = llm_models_to_choose, value = 'mistralai/Mistral-7B-Instruct-v0.2', label = 'LLM Model')
           llm_results = gr.Slider(minimum=4, maximum=10, value=5, step=1, interactive=True, label="Top n results as context")
-          stream_results = gr.Checkbox(value = True, label = "Stream output")
+          database_src = gr.Dropdown(choices = database_choices, value = index_info, label = 'Search Source')
+          stream_results = gr.Checkbox(value = True, label = "Stream output", visible = False)
 
     output_text = gr.Textbox(show_label = True, container = True, label = 'LLM Answer', visible = True, placeholder = output_placeholder)
     input = gr.Textbox(show_label = False, visible = False)
     gr_md = gr.Markdown(mark_text + md_text_initial)
 
-    def update_with_rag_md(message, llm_results_use = 5, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
-        rag_out = get_rag(message)
+    def update_with_rag_md(message, llm_results_use = 5, database_choice = index_info, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
+        prompt_text_from_data = ""
+        database_to_use = database_choice
+        if database_choice == index_info:
+          rag_out = get_rag(message)
+        else:
+          arxiv_search_success = True
+          try:
+            rag_out = get_arxiv_live_search(message, arx_client, retrieve_results)
+            if len(rag_out) == 0:
+              arxiv_search_success = False 
+          except:
+            arxiv_search_success = False
+ 
+
+          if not arxiv_search_success:
+            gr.Warning("Arxiv Search not working, switching to semantic search ...")
+            rag_out = get_rag(message)
+            database_to_use = index_info 
+
         md_text_updated = mark_text
         for i in range(retrieve_results):
           rag_answer = rag_out[i]
-          title = rag_answer['document_metadata']['title'].replace('\n','')
-
-          date = rag_answer['document_metadata']['_time']
-          paper_title = f'''### {date} | [{title}](https://arxiv.org/abs/{rag_answer['document_id']}) | [⬇️](https://arxiv.org/pdf/{rag_answer['document_id']})\n'''
-          paper_abs = rag_answer['content']
-          authors = rag_answer['document_metadata']['authors'].replace('\n','')
-          authors_formatted = f'*{authors}*' + ' \n\n'
-        
-          md_text_updated += paper_title + authors_formatted + paper_abs +  '\n---------------\n'+ '\n'
-        prompt = get_prompt_text(message, '\n\n'.join(rag_cleaner(out) for out in rag_out[:llm_results_use]), llm_model_picked = llm_model_picked)
+          if i < llm_results_use:
+            md_text_paper, prompt_text = get_md_text_abstract(rag_answer, source = database_to_use, return_prompt_formatting = True)
+            prompt_text_from_data += f"{i+1}. {prompt_text}"
+          else:
+            md_text_paper = get_md_text_abstract(rag_answer, source = database_to_use)
+          md_text_updated += md_text_paper
+        prompt = get_prompt_text(message, prompt_text_from_data, llm_model_picked = llm_model_picked)
         return md_text_updated, prompt
 
     def ask_llm(prompt, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2', stream_outputs = False):
@@ -144,6 +179,6 @@ with gr.Blocks(theme = gr.themes.Soft()) as demo:
            return stream
 
 
-    msg.submit(update_with_rag_md, [msg, llm_results, llm_model], [gr_md, input]).success(ask_llm, [input, llm_model, stream_results], output_text)
+    msg.submit(update_with_rag_md, [msg, llm_results,  database_src, llm_model], [gr_md, input]).success(ask_llm, [input, llm_model, stream_results], output_text)
 
 demo.queue().launch()

colbert/indexes/arxiv_colbert/0.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:02b786be4899edd43c88a271b10d650c702685b0de0198fd35d11cf776711d12
+oid sha256:9233ec6d9d53438a7d87ac1427e363ce0021b9b5856d14ee02b05ec420482a34
 size 17224796

colbert/indexes/arxiv_colbert/0.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:596b8f4f2a014c4f32e404fd71ef9f06b68501fa6c978f09ea09256a212999ca
+oid sha256:2a4fb10b5e65af293fd55722c825a34711e80173b928d56af051af2f7cf2c2a2
 size 137790512

colbert/indexes/arxiv_colbert/1.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5f032bc6b476431e756f4ebbfbeb7fe41c50e3ae5c00f8d4ebdf2ec89ba57966
+oid sha256:00ebcc4851d6d9ebe0773b49edb01229a9b64ce01b5d21c7598d00166f13be18
 size 18657244

colbert/indexes/arxiv_colbert/1.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8e35a3ef5f4327be903c8ba7d8e36321326b6372f7afc3e99f5dbc8851cd8253
+oid sha256:19e66224c0cd2080c5ae292121711c73e7b4b364c364fe32ac13b25eadc9a80c
 size 149249776

colbert/indexes/arxiv_colbert/10.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:07855f999051c8cf69c84e70e5c98250cd36fed529bff2051aca48daa66ed21e
-size 21300641
+oid sha256:90cd30a8d65798482e301f84e5588129814b66dd3c03a41e1fecbe62261d530e
+size 21302241

colbert/indexes/arxiv_colbert/10.metadata.json

@@ -1,6 +1,6 @@
 {
     "passage_offset": 250000,
     "num_passages": 25000,
-    "num_embeddings": 5324869,
-    "embedding_offset": 49558824
+    "num_embeddings": 5325266,
+    "embedding_offset": 49559057
 }

colbert/indexes/arxiv_colbert/10.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f95236110cc436b548b255001cf1fd18b707bd2098db6cd77b4c5749e7b6c78c
-size 170396981
+oid sha256:4e9425a008bfdf1583cf39a17822ecc3d9bfb5a23de64667bcd8c8dbe7399887
+size 170409717

colbert/indexes/arxiv_colbert/11.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:df3a9605cf26f8a94310238d4b3953ae915a7514732c1b2260b13b659e076ee3
-size 21343585
+oid sha256:6f62a9a28737ccdcfe088d3e8cab11a347b8a4bc61ea9fe4614fb411b0c3c3aa
+size 21343713

colbert/indexes/arxiv_colbert/11.metadata.json

@@ -1,6 +1,6 @@
 {
     "passage_offset": 275000,
     "num_passages": 25000,
-    "num_embeddings": 5335612,
-    "embedding_offset": 54883693
+    "num_embeddings": 5335645,
+    "embedding_offset": 54884323
 }

colbert/indexes/arxiv_colbert/11.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:10f19e52e503c4cbc514d0cc039b6700477bc1c986b26c6f5625978feda90353
-size 170740789
+oid sha256:e25eee695d53dab6bc3e7b02abb63cd82ed14297d270119cc785beb825b672b5
+size 170741813

colbert/indexes/arxiv_colbert/12.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8e740d1ae131a2b18a355017a2c464f53e927750a299176eeb05281f236752de
-size 21517665
+oid sha256:241ed1eba6c00d899e24195658138dd4706621cffa4e93366fa023ca18de3f24
+size 21518177

colbert/indexes/arxiv_colbert/12.metadata.json

@@ -1,6 +1,6 @@
 {
     "passage_offset": 300000,
     "num_passages": 25000,
-    "num_embeddings": 5379134,
-    "embedding_offset": 60219305
+    "num_embeddings": 5379263,
+    "embedding_offset": 60219968
 }

colbert/indexes/arxiv_colbert/12.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c6cf0ab13a172649ebac79caff4904b95f381ce1e92d9e7f983cda10005e1df6
-size 172133493
+oid sha256:8d252697525846a8a3d04375376b60aa8249f5848a48810354694a7befbde47c
+size 172137589

colbert/indexes/arxiv_colbert/13.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3b993050459a61fb92c0dc83c13eca96b579baa076032dab5c96bd3f20cacc5c
-size 14703969
+oid sha256:da50be3ade8c22b1d885e21c0255862c8b35d785a847f015caeb267d361429d6
+size 16133153

colbert/indexes/arxiv_colbert/13.metadata.json

@@ -1,6 +1,6 @@
 {
     "passage_offset": 325000,
-    "num_passages": 18284,
-    "num_embeddings": 3675697,
-    "embedding_offset": 65598439
+    "num_passages": 19957,
+    "num_embeddings": 4033002,
+    "embedding_offset": 65599231
 }

colbert/indexes/arxiv_colbert/13.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:74abf66bdbb8b8ca205bbf28ecd62c5c99fea5c9ffccdc86684c019e032c2acc
-size 117623477
+oid sha256:5ddb758cc21cae775847adf144dc81f49d6deaa328cbb3a4b4decfc8dc55da85
+size 129057269

colbert/indexes/arxiv_colbert/2.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cb9c97c10f869355677d800cb3d1480991522736005813635f09a1657250b7d5
-size 19123036
+oid sha256:f2a37f8390b58ace74df2176722ae65c31f1ba3986bc69f90b233d3fc8c32491
+size 19122908

colbert/indexes/arxiv_colbert/2.metadata.json

@@ -1,6 +1,6 @@
 {
     "passage_offset": 50000,
     "num_passages": 25000,
-    "num_embeddings": 4780476,
+    "num_embeddings": 4780440,
     "embedding_offset": 8969935
 }

colbert/indexes/arxiv_colbert/2.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:04c15e645fbcee5378246466d9fc206b491f3e0daad7e58d44682a450cb419c4
-size 152976432
+oid sha256:9ea1f7004b9dd497306bdca872fc181e803a937203253d38081a62ad27f2b2af
+size 152975280

colbert/indexes/arxiv_colbert/3.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1c3b7b142607749941de818debc51a9dc14757ad069c84e3dc796f644202f399
+oid sha256:127a154192c8444def0e5f11b9ffb845e60947513cff266b6d399bb8e0939a6f
 size 19531484

colbert/indexes/arxiv_colbert/3.metadata.json

@@ -2,5 +2,5 @@
     "passage_offset": 75000,
     "num_passages": 25000,
     "num_embeddings": 4882584,
-    "embedding_offset": 13750411
+    "embedding_offset": 13750375
 }

colbert/indexes/arxiv_colbert/3.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:eaaff7e6e361f26f8be33201d391d83f6a58491d8d236f8531e46e430a96c857
+oid sha256:e98eed69ae9b86900bed76837111982a91f80ba51be8d3ce0e10477df7cdaf67
 size 156243888

colbert/indexes/arxiv_colbert/4.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:26b7424ec40d6aee6cd938daf31f0c6769b76e9c57f2d83f0d3876b0f331b524
+oid sha256:647e2266da43a402faac17f51c870b96b7c1e915c33758dc3cbf6bf6edbf96ae
 size 19962524

colbert/indexes/arxiv_colbert/4.metadata.json

@@ -1,6 +1,6 @@
 {
     "passage_offset": 100000,
     "num_passages": 25000,
-    "num_embeddings": 4990348,
-    "embedding_offset": 18632995
+    "num_embeddings": 4990349,
+    "embedding_offset": 18632959
 }

colbert/indexes/arxiv_colbert/4.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9c2b8a4b85ba90d2c4fc76bb6d0c75cf7222ee878088dd80d1d8a41f0560c612
+oid sha256:e428585dc9a4611abe8e789c3f34ac992d2f7398e9203b183a06d14179c07431
 size 159692336

colbert/indexes/arxiv_colbert/5.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e6e121cf813e9c01dd1db42391f3ded109d2ff23eef42c0cfd2291a561ba3b15
-size 20303068
+oid sha256:e6d3754ba1252b5943ff2502d84dc5857d43b281a4decab589a58d5052ab44e2
+size 20303580

colbert/indexes/arxiv_colbert/5.metadata.json

@@ -1,6 +1,6 @@
 {
     "passage_offset": 125000,
     "num_passages": 25000,
-    "num_embeddings": 5075480,
-    "embedding_offset": 23623343
+    "num_embeddings": 5075614,
+    "embedding_offset": 23623308
 }

colbert/indexes/arxiv_colbert/5.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fe127677d3ac6f547045bc4e08c7d02f342ee0fb74b9930e28a4fb0def30906e
-size 162416560
+oid sha256:84115a366b7cb6aaf6f971a916d064ac73e878a9cb63960190268e78807c3b7c
+size 162420848

colbert/indexes/arxiv_colbert/6.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7b4596a5628764e7a81c77e77831be5a63d18870c2060a311b30856e240b457b
-size 20472796
+oid sha256:9c2e2c32c99abc9191959eb90fdc3e251fca517c78efeddc499d402640e5fb2c
+size 20473180

colbert/indexes/arxiv_colbert/6.metadata.json

@@ -1,6 +1,6 @@
 {
     "passage_offset": 150000,
     "num_passages": 25000,
-    "num_embeddings": 5117915,
-    "embedding_offset": 28698823
+    "num_embeddings": 5118011,
+    "embedding_offset": 28698922
 }

colbert/indexes/arxiv_colbert/6.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e1fa2fd34658f6099bdf17f5154c054adae2f6969e646d57a6785165f0739c8f
-size 163774448
+oid sha256:36c449b807be2ee56d03a1a1d3bace6a7702b3682715f011c682712cd172b120
+size 163777520

colbert/indexes/arxiv_colbert/7.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fd442e5168eedc3762451461ae72c680d870e70a2e3a74328541d69137e14eb2
-size 20836892
+oid sha256:c03f4544c82f93eda118eb8a4cce6bf9636a4bc05cb248b4bcbe3b4c08dc3e3f
+size 20836636

colbert/indexes/arxiv_colbert/7.metadata.json

@@ -1,6 +1,6 @@
 {
     "passage_offset": 175000,
     "num_passages": 25000,
-    "num_embeddings": 5208942,
-    "embedding_offset": 33816738
+    "num_embeddings": 5208871,
+    "embedding_offset": 33816933
 }

colbert/indexes/arxiv_colbert/7.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:089557df32c0678b171d6fa8cd077032be52ec1362952a2e402157edd39f12e4
-size 166687344
+oid sha256:2692e7da6ced1b077c1338e2725085c140c84c6868e31a8c28a7924b7383cd57
+size 166685040

colbert/indexes/arxiv_colbert/8.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:180f6b3a35464ca0c2911df592d817271afa30695c6fcc42745f7db24a6514fc
-size 21021340
+oid sha256:d621ba8459bf75fedc6c3800c57e01f6c94ecd067396c3a9f5dbcbd851a100c2
+size 21021404

colbert/indexes/arxiv_colbert/8.metadata.json

@@ -1,6 +1,6 @@
 {
     "passage_offset": 200000,
     "num_passages": 25000,
-    "num_embeddings": 5255043,
-    "embedding_offset": 39025680
+    "num_embeddings": 5255060,
+    "embedding_offset": 39025804
 }

colbert/indexes/arxiv_colbert/8.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8d3b5159f26fec80325a56228f945d92bd15b27e1d3a704a2acd992075c0e71a
-size 168162544
+oid sha256:d371e1d450e2092cc841c9aa2a3eba8927cc2e2f98867f52ae183a87d27247df
+size 168163120

colbert/indexes/arxiv_colbert/9.codes.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:be0ccfd90a5385296aa821d90f9ff15c78d7ef2ea2967adc14a1a75a171cb69a
-size 21113564
+oid sha256:a94be0d3ccf5606f49ca0c11f1007a028b14faaaeb77730ad65b1877f26a5e4f
+size 21113948

colbert/indexes/arxiv_colbert/9.metadata.json

@@ -1,6 +1,6 @@
 {
     "passage_offset": 225000,
     "num_passages": 25000,
-    "num_embeddings": 5278101,
-    "embedding_offset": 44280723
+    "num_embeddings": 5278193,
+    "embedding_offset": 44280864
 }

colbert/indexes/arxiv_colbert/9.residuals.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e4a7d456fa97a60a90d0bb1c8d6caa05033364f6a5b75917ffd6bf7a85e3717a
-size 168900400
+oid sha256:04f3d5d2d05dceb31430cdcaf7f92151140e72d05a669c463e2354fcea65c5b3
+size 168903344

colbert/indexes/arxiv_colbert/buckets.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:804931d6454d91dbcd74a19b2345189aa49b5056fce6582f7ede479a9730f52c
+oid sha256:17a58ea162ec2bbbc1dd21141742cbc82dbfe3a13d1d14c44f64e1603ee27b51
 size 1432

colbert/indexes/arxiv_colbert/centroids.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2bc89b1ea0d60711634d5ec75d28580a750a7e073732157bec3b5319456ac401
+oid sha256:0f8cc9ba855afe896e60ababdb0296303783440ba6505f6c9d533fd2bb165e2d
 size 33555622

colbert/indexes/arxiv_colbert/collection.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bbda587dd17f68c76673cd326294004eee932d542fd927c043789fb61db65bba
-size 409810003
+oid sha256:666e887d3dc9aa45fa575c60547a4d151822e840feca87f48d17740ac0891f89
+size 411963323

colbert/indexes/arxiv_colbert/docid_metadata_map.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a61ea67a0fb6939795f31877e69713cb83db43ec64148726841bedcb5816d9c0
-size 84148235
+oid sha256:76b05799713b16935bf43c6ac1f45bd466803c2cbc4201c367b09bc2ac0f4653
+size 84578105