Add model and some other thing

app.py

@@ -9,10 +9,12 @@ from time import perf_counter
 import gradio as gr
 from jinja2 import Environment, FileSystemLoader
 
-from backend.query_llm import generate_hf, generate_openai
+from backend.query_llm import generate_hf, generate_openai, hf_models, openai_models
 from backend.semantic_search import retrieve
 import itertools
 
+inf_models = list(hf_models.keys()) + list(openai_models)
+
 emb_models = ["bge", "minilm"]
 splitters = ['ct', 'rct', 'nltk']
 chunk_sizes = ["500", "2000"]
@@ -56,7 +58,8 @@ def has_balanced_backticks(markdown_str):
     # If in_code_block is False at the end, all backticks are balanced
     return not in_code_block
 
-def bot(history, api_kind,
+def bot(history, model_name, oepnai_api_key,
+                  reranker_enabled,reranker_kind,num_prerank_docs,
             num_docs, model_kind, sub_vector_size, chunk_size, splitter_type, all_at_once):
     query = history[-1][0]
 
@@ -67,7 +70,10 @@ def bot(history, api_kind,
     # Retrieve documents relevant to query
     document_start = perf_counter()
 
-    documents = retrieve(query, int(num_docs), model_kind, sub_vector_size, chunk_size, splitter_type)
+    if reranker_enabled:
+        documents = retrieve(query, int(num_docs), model_kind, sub_vector_size, chunk_size, splitter_type,reranker_kind,num_prerank_docs)
+    else:
+        documents = retrieve(query, int(num_docs), model_kind, sub_vector_size, chunk_size, splitter_type)
 
     document_time = perf_counter() - document_start
     logger.info(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
@@ -77,18 +83,18 @@ def bot(history, api_kind,
     prompt_html = template_html.render(documents=documents, query=query)
 
 
-    if api_kind == "HuggingFace":
+    if model_name == "HuggingFace":
          generate_fn = generate_hf
-    elif api_kind == "OpenAI":
+    elif model_name == "OpenAI":
          generate_fn = generate_openai
     else:
-         raise gr.Error(f"API {api_kind} is not supported")
+         raise gr.Error(f"API {model_name} is not supported")
     
 
     history[-1][1] = ""
     if all_at_once:
-        for model_name, doc, size, sub_vector in combinations:
-            documents_i = retrieve(query, int(num_docs), model_name, sub_vector, size, doc)
+        for emb_model, doc, size, sub_vector in combinations:
+            documents_i = retrieve(query, int(num_docs), emb_model, sub_vector, size, doc)
             prompt_i = template.render(documents=documents_i, query=query)
             prompt_html = template_html.render(documents=documents, query=query)
             
@@ -96,13 +102,13 @@ def bot(history, api_kind,
             prev_hist = history[-1][1] 
             if not has_balanced_backticks(prev_hist):
                 prev_hist += "\n```\n"
-            prev_hist += f"\n\n## model {model_name}, splitter {doc}, size {size}, sub vector {sub_vector}\n\n"
-            for character in generate_fn(prompt_i, history[:-1]):
+            prev_hist += f"\n\n## model {emb_model}, splitter {doc}, size {size}, sub vector {sub_vector}\n\n"
+            for character in generate_fn(model_name, prompt_i, history[:-1], oepnai_api_key):
                 hist_chunk = character
                 history[-1][1] = prev_hist + hist_chunk
                 yield history, prompt_html
     else:
-        for character in generate_fn(prompt, history[:-1]):
+        for character in generate_fn(model_name, prompt, history[:-1], oepnai_api_key):
             history[-1][1] = character
             yield history, prompt_html
     
@@ -129,20 +135,30 @@ with gr.Blocks() as demo:
 
 
     with gr.Row():
-        num_docs = gr.Slider(1, 20, label="number of docs", step=1, value=4)
-        model_kind = gr.Radio(choices=emb_models, value="bge", label="embedding model")
+        emb_model_kind = gr.Radio(choices=emb_models, value="bge", label="embedding model")
         sub_vector_size = gr.Radio(choices=sub_vectors, value="32", label="sub-vector size")
-        all_at_once = gr.Checkbox(value=False, label="Run all at once")
-    with gr.Row():
-        api_kind = gr.Radio(choices=["HuggingFace", "OpenAI"], value="HuggingFace", label="Chat model engine")
         chunk_size = gr.Radio(choices=chunk_sizes, value="2000", label="chunk size")
         splitter_type = gr.Radio(choices=splitters, value="nltk", label="splitter")
+    with gr.Row():
+        reranker_enabled = gr.Checkbox(value=False, label="Reranker enabled")
+        reranker_kind = gr.Radio(choices=emb_models, value="bge", label="Reranker model")
+        num_prerank_docs = gr.Slider(5, 80, label="Number of docs before reranker", step=1, value=20)
+    with gr.Row():
+        num_docs = gr.Slider(1, 20, label="number of docs", step=1, value=4)
+        all_at_once = gr.Checkbox(value=False, label="Run all at once")
+        model_name = gr.Radio(choices=inf_models, value=inf_models[0], label="Chat model")
+        oepnai_api_key = gr.Textbox(
+                show_label=False,
+                placeholder="OpenAI API key",
+                container=False,
+                )
 
     prompt_html = gr.HTML()
     # Turn off interactivity while generating if you click
     txt_msg = txt_btn.click(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-            bot, [chatbot, api_kind,
-            num_docs, model_kind, sub_vector_size, chunk_size, splitter_type, all_at_once
+            bot, [chatbot, model_name, oepnai_api_key,
+                  reranker_enabled,reranker_kind,num_prerank_docs,
+                num_docs, emb_model_kind, sub_vector_size, chunk_size, splitter_type, all_at_once
             ], [chatbot, prompt_html])
 
     # Turn it back on
@@ -150,8 +166,8 @@ with gr.Blocks() as demo:
 
     # Turn off interactivity while generating if you hit enter
     txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-            bot, [chatbot, api_kind,
-            num_docs, model_kind, sub_vector_size, chunk_size, splitter_type
+            bot, [chatbot, model_name,
+            num_docs, emb_model_kind, sub_vector_size, chunk_size, splitter_type
             ], [chatbot, prompt_html])
 
     # Turn it back on

backend/query_llm.py

@@ -8,15 +8,19 @@ from huggingface_hub import InferenceClient
 from transformers import AutoTokenizer
 
 
-OPENAI_KEY = os.getenv("OPENAI_API_KEY")
 HF_TOKEN = os.getenv("HF_TOKEN")
-TOKENIZER = AutoTokenizer.from_pretrained(os.getenv("HF_MODEL"))
 
-HF_CLIENT = InferenceClient(
-    os.getenv("HF_MODEL"),
-    token=HF_TOKEN
-)
-OAI_CLIENT = openai.Client(api_key=OPENAI_KEY)
+hf_models = {
+    "mistral-7B": "mistralai/Mistral-7B-Instruct-v0.2",
+    "mistral-7B 0.1": "mistralai/Mistral-7B-v0.1",
+    "llama 3": "meta-llama/Meta-Llama-3-70B-Instruct",
+}
+openai_models = {"gpt-4o","gpt-3.5-turbo-0125"}
+
+tokenizers = {k: AutoTokenizer.from_pretrained(m) for k,m in hf_models.items()}
+clients = {k: InferenceClient(
+    m, token=HF_TOKEN
+) for k,m in hf_models.items()}
 
 HF_GENERATE_KWARGS = {
     'temperature': max(float(os.getenv("TEMPERATURE", 0.9)), 1e-2),
@@ -34,7 +38,7 @@ OAI_GENERATE_KWARGS = {
 }
 
 
-def format_prompt(message: str, api_kind: str):
+def format_prompt(message: str, model: str):
     """
     Formats the given message using a chat template.
 
@@ -48,15 +52,15 @@ def format_prompt(message: str, api_kind: str):
     # Create a list of message dictionaries with role and content
     messages: List[Dict[str, Any]] = [{'role': 'user', 'content': message}]
 
-    if api_kind == "openai":
+    if model in openai_models:
         return messages
-    elif api_kind == "hf":
-        return TOKENIZER.apply_chat_template(messages, tokenize=False)
-    elif api_kind:
-        raise ValueError("API is not supported")
+    elif model in hf_models:
+        return tokenizers[model].apply_chat_template(messages, tokenize=False)
+    else:
+        raise ValueError(f"Model {model} is not supported")
 
 
-def generate_hf(prompt: str, history: str) -> Generator[str, None, str]:
+def generate_hf(model: str, prompt: str, history: str, _: str) -> Generator[str, None, str]:
     """
     Generate a sequence of tokens based on a given prompt and history using Mistral client.
 
@@ -68,11 +72,11 @@ def generate_hf(prompt: str, history: str) -> Generator[str, None, str]:
                                    Returns a final string if an error occurs.
     """
 
-    formatted_prompt = format_prompt(prompt, "hf")
+    formatted_prompt = format_prompt(prompt, model)
     formatted_prompt = formatted_prompt.encode("utf-8").decode("utf-8")
 
     try:
-        stream = HF_CLIENT.text_generation(
+        stream = clients[model].text_generation(
             formatted_prompt,
             **HF_GENERATE_KWARGS,
             stream=True,
@@ -93,7 +97,7 @@ def generate_hf(prompt: str, history: str) -> Generator[str, None, str]:
             raise gr.Error(f"Unhandled Exception: {str(e)}")
 
 
-def generate_openai(prompt: str, history: str) -> Generator[str, None, str]:
+def generate_openai(model: str, prompt: str, history: str, api_key: str) -> Generator[str, None, str]:
     """
     Generate a sequence of tokens based on a given prompt and history using Mistral client.
 
@@ -104,11 +108,12 @@ def generate_openai(prompt: str, history: str) -> Generator[str, None, str]:
         Generator[str, None, str]: A generator yielding chunks of generated text.
                                    Returns a final string if an error occurs.
     """
-    formatted_prompt = format_prompt(prompt, "openai")
+    formatted_prompt = format_prompt(prompt, model)
+    client = openai.Client(api_key=api_key)
 
     try:
-        stream = OAI_CLIENT.chat.completions.create(
-            model=os.getenv("OPENAI_MODEL"),
+        stream = client.chat.completions.create(
+            model=model,
             messages=formatted_prompt,
             **OAI_GENERATE_KWARGS,
             stream=True

backend/semantic_search.py

@@ -1,7 +1,7 @@
 import lancedb
 import os
 import gradio as gr
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, CrossEncoder
 
 
 db = lancedb.connect(".lancedb")
@@ -22,20 +22,41 @@ BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
 retriever_bge = SentenceTransformer("BAAI/bge-large-en-v1.5")
 retriever_minilm = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
+reranker_bge = CrossEncoder("BAAI/bge-reranker-large")
+reranker_minilm = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
 
-def retrieve(query, k, model_kind, sub_vector_size, chunk_size, splitter_type):
+
+def retrieve(query, k, model_kind, sub_vector_size, chunk_size, splitter_type, reranker_kind=None, pre_ranker_size=None):
     if model_kind == "bge":
         query_vec = retriever_bge.encode(query)
     else:
         query_vec = retriever_minilm.encode(query)
 
+    if pre_ranker_size is None:
+        pre_ranker_size = k
+
     try:
         documents = table(
             f"{splitter_type}_{model_kind}_{sub_vector_size}_{chunk_size}",
-            ).search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k).to_list()
+            ).search(query_vec, vector_column_name=VECTOR_COLUMN).limit(pre_ranker_size).to_list()
         documents = [doc[TEXT_COLUMN] for doc in documents]
 
-        return documents
+        if reranker_kind is None:
+            return documents
+        # Pair the query with each document for re-ranking
+        query_document_pairs = [(query, text) for text in documents]
+
+        # Score documents using the reranker
+        if reranker_kind == "bge":
+            scores = reranker_bge.predict(query_document_pairs)
+        else:
+            scores = reranker_minilm.predict(query_document_pairs)
+
+        # Aggregate and sort the documents based on the scores
+        scored_documents = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
+
+        # Return the top K documents based on re-ranking
+        return [doc for doc, _ in scored_documents[:k]]
 
     except Exception as e:
         raise gr.Error(str(e))