Update app.py

app.py

@@ -2,7 +2,7 @@
 """
 Credit to Derek Thomas, [email protected]
 """
-
+from ragatouille import RAGPretrainedModel
 import subprocess
 
 # subprocess.run(["pip", "install", "--upgrade", "transformers[torch,sentencepiece]==4.34.1"])
@@ -59,46 +59,68 @@ def bot(history, cross_encoder):
          raise ValueError("Empty string was submitted")
 
     logger.warning('Retrieving documents...')
-    # Retrieve documents relevant to query
-    document_start = perf_counter()
-
-    query_vec = retriever.encode(query)
-    logger.warning(f'Finished query vec')
-    doc1 = table.search(query_vec, vector_column_name=VECTOR_COLUMN_NAME).limit(top_k_rank)
-
     
-
-    logger.warning(f'Finished search')
-    documents = table.search(query_vec, vector_column_name=VECTOR_COLUMN_NAME).limit(top_rerank).to_list()
-    documents = [doc[TEXT_COLUMN_NAME] for doc in documents]
-    logger.warning(f'start cross encoder {len(documents)}')
-    # Retrieve documents relevant to query
-    query_doc_pair = [[query, doc] for doc in documents]
-    if cross_encoder=='MiniLM-L6v2' :
-           cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') 
+    # if COLBERT RAGATATOUILLE PROCEDURE  : 
+    if cross_encoder=='ColBERT':
+        gr.Warning('Retrieving using ColBERT')
+        RAG= RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")
+        RAG_db=RAG.from_index('.ragatouille/colbert/indexes/mockingbird')
+        documents_full=RAG_db.search(query)
+        
+        documents=[item['content'] for item in documents_full]
+        # Create Prompt
+        prompt = template.render(documents=documents, query=query)
+        prompt_html = template_html.render(documents=documents, query=query)
+    
+        generate_fn = generate_hf
+    
+        history[-1][1] = ""
+        for character in generate_fn(prompt, history[:-1]):
+            history[-1][1] = character
+            print('Final history is ',history)
+            yield history, prompt_html
     else:
-           cross_encoder = CrossEncoder('BAAI/bge-reranker-base')
-    cross_scores = cross_encoder.predict(query_doc_pair)
-    sim_scores_argsort = list(reversed(np.argsort(cross_scores)))
-    logger.warning(f'Finished cross encoder {len(documents)}')
+        # Retrieve documents relevant to query
+        document_start = perf_counter()
     
-    documents = [documents[idx] for idx in sim_scores_argsort[:top_k_rank]]
-    logger.warning(f'num documents {len(documents)}')
-
-    document_time = perf_counter() - document_start
-    logger.warning(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
-
-    # Create Prompt
-    prompt = template.render(documents=documents, query=query)
-    prompt_html = template_html.render(documents=documents, query=query)
-
-    generate_fn = generate_hf
-
-    history[-1][1] = ""
-    for character in generate_fn(prompt, history[:-1]):
-        history[-1][1] = character
-        print('Final history is ',history)
-        yield history, prompt_html
+        query_vec = retriever.encode(query)
+        logger.warning(f'Finished query vec')
+        doc1 = table.search(query_vec, vector_column_name=VECTOR_COLUMN_NAME).limit(top_k_rank)
+    
+        
+    
+        logger.warning(f'Finished search')
+        documents = table.search(query_vec, vector_column_name=VECTOR_COLUMN_NAME).limit(top_rerank).to_list()
+        documents = [doc[TEXT_COLUMN_NAME] for doc in documents]
+        logger.warning(f'start cross encoder {len(documents)}')
+        # Retrieve documents relevant to query
+        query_doc_pair = [[query, doc] for doc in documents]
+        if cross_encoder=='MiniLM-L6v2' :
+               cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') 
+        elif cross_encoder=='BGE reranker':
+               cross_encoder = CrossEncoder('BAAI/bge-reranker-base')
+        
+        cross_scores = cross_encoder.predict(query_doc_pair)
+        sim_scores_argsort = list(reversed(np.argsort(cross_scores)))
+        logger.warning(f'Finished cross encoder {len(documents)}')
+        
+        documents = [documents[idx] for idx in sim_scores_argsort[:top_k_rank]]
+        logger.warning(f'num documents {len(documents)}')
+    
+        document_time = perf_counter() - document_start
+        logger.warning(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
+    
+        # Create Prompt
+        prompt = template.render(documents=documents, query=query)
+        prompt_html = template_html.render(documents=documents, query=query)
+    
+        generate_fn = generate_hf
+    
+        history[-1][1] = ""
+        for character in generate_fn(prompt, history[:-1]):
+            history[-1][1] = character
+            print('Final history is ',history)
+            yield history, prompt_html
 
 
 with gr.Blocks(theme='Insuz/SimpleIndigo') as demo:
@@ -128,7 +150,7 @@ with gr.Blocks(theme='Insuz/SimpleIndigo') as demo:
                 )
         txt_btn = gr.Button(value="Submit text", scale=1)
 
-    cross_encoder = gr.Radio(choices=['MiniLM-L6v2','BGE reranker'], value='BGE reranker',label="Embeddings", info="Choose MiniLM for Speed, BGE reranker for accuracy")
+    cross_encoder = gr.Radio(choices=['MiniLM-L6v2','BGE reranker','ColBERT'], value='BGE reranker',label="Embeddings", info="Choose MiniLM for Speed, BGE reranker for accuracy,ColBERT for both")
 
     prompt_html = gr.HTML()
     # Turn off interactivity while generating if you click