Add GPT retrieval

app.py

@@ -2,8 +2,9 @@ import pandas as pd
 import gradio as gr
 from transformers import pipeline
 import nltk
-from retrieval import retrieve_from_pdf
-from llm_retrieval import retrieve_from_pdf_llm, retrieve_from_pdf_llm_fast
+from retrieval_bm25s import retrieve_with_bm25s
+from retrieval_bert import retrieve_with_deberta
+from retrieval_gpt import retrieve_with_gpt
 import os
 import json
 from datetime import datetime
@@ -102,10 +103,10 @@ with gr.Blocks(theme=my_theme, head=font_awesome_html) as demo:
                         )
                         with gr.Row():
                             retrieval_method = gr.Radio(
-                                choices=["BM25S", "LLM (Large)", "LLM (Fast)"],
+                                choices=["BM25S", "DeBERTa", "GPT"],
                                 value="BM25S",
                                 label="Retrieval Method",
-                                info="Choose between keyword-based (BM25S) or AI-based (LLM) evidence retrieval",
+                                info="Keyword search (BM25S) or AI (DeBERTa, GPT)",
                             )
                         get_evidence = gr.Button(value="Get Evidence")
                         top_k = gr.Slider(
@@ -113,7 +114,6 @@ with gr.Blocks(theme=my_theme, head=font_awesome_html) as demo:
                             10,
                             value=5,
                             step=1,
-                            interactive=True,
                             label="Top k sentences",
                         )
                 with gr.Column(scale=3):
@@ -122,7 +122,11 @@ with gr.Blocks(theme=my_theme, head=font_awesome_html) as demo:
                         info="aka premise",
                         placeholder="Input evidence or use Get Evidence from PDF",
                     )
-            submit = gr.Button("3. Submit", visible=False)
+                    with gr.Row():
+                        prompt_tokens = gr.Number(label="Prompt tokens", visible=False)
+                        completion_tokens = gr.Number(
+                            label="Completion tokens", visible=False
+                        )
 
         with gr.Column(scale=2):
             # Keep the prediction textbox hidden
@@ -234,17 +238,17 @@ with gr.Blocks(theme=my_theme, head=font_awesome_html) as demo:
                 #### *Capstone project*
                 - <i class="fa-brands fa-github"></i> [jedick/MLE-capstone-project](https://github.com/jedick/MLE-capstone-project) (project repo)
                 - <i class="fa-brands fa-github"></i> [jedick/AI4citations](https://github.com/jedick/AI4citations) (app repo)
-                #### *Claim Verification Models (text classification)*
+                #### *Text Classification*
                 - <img src="https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.svg" style="height: 1.2em; display: inline-block;"> [jedick/DeBERTa-v3-base-mnli-fever-anli-scifact-citint](https://huggingface.co/jedick/DeBERTa-v3-base-mnli-fever-anli-scifact-citint) (fine-tuned)
                 - <img src="https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.svg" style="height: 1.2em; display: inline-block;"> [MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli](https://huggingface.co/MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli) (base)
-                #### *Evidence Retrieval Models (question answering)*
-                - <img src="https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.svg" style="height: 1.2em; display: inline-block;"> [deepset/deberta-v3-large-squad2](https://huggingface.co/deepset/deberta-v3-large-squad2) (Large)
-                - <img src="https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.svg" style="height: 1.2em; display: inline-block;"> [distilbert-base-cased-distilled-squad](https://huggingface.co/distilbert/distilbert-base-cased-distilled-squad) (Fast)
+                #### *Evidence Retrieval*
+                - <i class="fa-brands fa-github"></i> [xhluca/bm25s](https://github.com/xhluca/bm25s) (BM25S)
+                - <img src="https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.svg" style="height: 1.2em; display: inline-block;"> [deepset/deberta-v3-large-squad2](https://huggingface.co/deepset/deberta-v3-large-squad2) (DeBERTa)
+                - <img src="https://upload.wikimedia.org/wikipedia/commons/4/4d/OpenAI_Logo.svg" style="height: 1.2em; display: inline-block;"> [gpt-4o-mini-2024-07-18](https://platform.openai.com/docs/pricing) (GPT)
                 #### *Datasets for fine-tuning*
                 - <i class="fa-brands fa-github"></i> [allenai/SciFact](https://github.com/allenai/scifact) (SciFact)
                 - <i class="fa-brands fa-github"></i> [ScienceNLP-Lab/Citation-Integrity](https://github.com/ScienceNLP-Lab/Citation-Integrity) (CitInt)
                 #### *Other sources*
-                - <i class="fa-brands fa-github"></i> [xhluca/bm25s](https://github.com/xhluca/bm25s) (evidence retrieval)
                 - <img src="https://plos.org/wp-content/uploads/2020/01/logo-color-blue.svg" style="height: 1.4em; display: inline-block;"> [Medicine](https://doi.org/10.1371/journal.pmed.0030197), <i class="fa-brands fa-wikipedia-w"></i> [CRISPR](https://en.wikipedia.org/wiki/CRISPR) (evidence retrieval examples)
                 - <img src="https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.svg" style="height: 1.2em; display: inline-block;"> [nyu-mll/multi_nli](https://huggingface.co/datasets/nyu-mll/multi_nli/viewer/default/train?row=37&views%5B%5D=train) (MNLI example)
                 - <img src="https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.svg" style="height: 1.2em; display: inline-block;"> [NoCrypt/miku](https://huggingface.co/spaces/NoCrypt/miku) (theme)
@@ -329,16 +333,23 @@ with gr.Blocks(theme=my_theme, head=font_awesome_html) as demo:
         return pdf_file, claim
 
     @spaces.GPU()
-    def retrieve_evidence_with_method(pdf_file, claim, top_k, method):
+    def _retrieve_with_deberta(pdf_file, claim, top_k):
+        """
+        Retrieve evidence using DeBERTa
+        """
+        return retrieve_with_deberta(pdf_file, claim, top_k)
+
+    def retrieve_evidence(pdf_file, claim, top_k, method):
         """
         Retrieve evidence using the selected method
         """
         if method == "BM25S":
-            return retrieve_from_pdf(pdf_file, claim, k=top_k)
-        elif method == "LLM (Large)":
-            return retrieve_from_pdf_llm(pdf_file, claim, k=top_k)
-        elif method == "LLM (Fast)":
-            return retrieve_from_pdf_llm_fast(pdf_file, claim, k=top_k)
+            # Append 0 for number of prompt and completion tokens
+            return retrieve_with_bm25s(pdf_file, claim, top_k), 0, 0
+        elif method == "DeBERTa":
+            return _retrieve_with_deberta(pdf_file, claim, top_k), 0, 0
+        elif method == "GPT":
+            return retrieve_with_gpt(pdf_file, claim)
         else:
             return f"Unknown retrieval method: {method}"
 
@@ -399,11 +410,29 @@ with gr.Blocks(theme=my_theme, head=font_awesome_html) as demo:
         else:
             append_feedback(*args, user_label="REFUTE")
 
+    def number_visible(value):
+        """
+        Show numbers (token counts) if GPT is selcted for retrieval
+        """
+        if value == "GPT":
+            return gr.Number(visible=True)
+        else:
+            return gr.Number(visible=False)
+
+    def slider_visible(value):
+        """
+        Hide slider (top_k) if GPT is selcted for retrieval
+        """
+        if value == "GPT":
+            return gr.Slider(visible=False)
+        else:
+            return gr.Slider(visible=True)
+
     # Event listeners
 
-    # Click the submit button or press Enter to submit
+    # Press Enter or Shift-Enter to submit
     gr.on(
-        triggers=[claim.submit, evidence.submit, submit.click],
+        triggers=[claim.submit, evidence.submit],
         fn=query_model,
         inputs=[claim, evidence],
         outputs=[prediction, label],
@@ -412,9 +441,9 @@ with gr.Blocks(theme=my_theme, head=font_awesome_html) as demo:
     # Get evidence from PDF and run the model
     gr.on(
         triggers=[get_evidence.click],
-        fn=retrieve_evidence_with_method,
+        fn=retrieve_evidence,
         inputs=[pdf_file, claim, top_k, retrieval_method],
-        outputs=evidence,
+        outputs=[evidence, prompt_tokens, completion_tokens],
     ).then(
         fn=query_model,
         inputs=[claim, evidence],
@@ -472,9 +501,9 @@ with gr.Blocks(theme=my_theme, head=font_awesome_html) as demo:
         outputs=[pdf_file, claim],
         api_name=False,
     ).then(
-        fn=retrieve_evidence_with_method,
+        fn=retrieve_evidence,
         inputs=[pdf_file, claim, top_k, retrieval_method],
-        outputs=evidence,
+        outputs=[evidence, prompt_tokens, completion_tokens],
         api_name=False,
     ).then(
         fn=query_model,
@@ -515,17 +544,29 @@ with gr.Blocks(theme=my_theme, head=font_awesome_html) as demo:
         fn=save_feedback_support,
         inputs=[claim, evidence, model, label],
         outputs=None,
+        api_name=False,
     )
     flag_nei.click(
         fn=save_feedback_nei,
         inputs=[claim, evidence, model, label],
         outputs=None,
+        api_name=False,
     )
     flag_refute.click(
         fn=save_feedback_refute,
         inputs=[claim, evidence, model, label],
         outputs=None,
+        api_name=False,
+    )
+
+    # Change visibility of top-k slider and token counts if GPT is selected for retrieval
+    retrieval_method.change(
+        number_visible, retrieval_method, prompt_tokens, api_name=False
+    )
+    retrieval_method.change(
+        number_visible, retrieval_method, completion_tokens, api_name=False
     )
+    retrieval_method.change(slider_visible, retrieval_method, top_k, api_name=False)
 
 
 if __name__ == "__main__":

requirements.txt

@@ -8,3 +8,4 @@ nltk
 bm25s
 huggingface_hub
 spaces
+openai

llm_retrieval.py → retrieval_bert.py

@@ -12,14 +12,14 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-class LLMEvidenceRetriever:
+class BERTRetriever:
     """
-    LLM-based evidence retrieval using extractive question answering
+    BERT-based evidence retrieval using extractive question answering
     """
 
     def __init__(self, model_name: str = "deepset/deberta-v3-large-squad2"):
         """
-        Initialize the LLM evidence retriever
+        Initialize the BERT evidence retriever
 
         Args:
             model_name: HuggingFace model for question answering
@@ -34,7 +34,7 @@ class LLMEvidenceRetriever:
         )
         # Maximum context length for the model
         self.max_length = self.tokenizer.model_max_length
-        logger.info(f"Initialized LLM retriever with model: {model_name}")
+        logger.info(f"Initialized BERT retriever with model: {model_name}")
 
     def _extract_and_clean_text(self, pdf_file: str) -> str:
         """
@@ -124,9 +124,9 @@ class LLMEvidenceRetriever:
         else:
             return f"What evidence supports the claim that {claim.lower()}?"
 
-    def retrieve_evidence(self, pdf_file: str, claim: str, k: int = 5) -> str:
+    def retrieve_evidence(self, pdf_file: str, claim: str, top_k: int = 5) -> str:
         """
-        Retrieve evidence from PDF using LLM-based question answering
+        Retrieve evidence from PDF using BERT-based question answering
 
         Args:
             pdf_file: Path to PDF file
@@ -177,7 +177,7 @@ class LLMEvidenceRetriever:
 
             # Sort by confidence score and take top k
             answers.sort(key=lambda x: x["score"], reverse=True)
-            top_answers = answers[:k]
+            top_answers = answers[:top_k]
 
             # Combine evidence passages
             if top_answers:
@@ -189,30 +189,30 @@ class LLMEvidenceRetriever:
                 return "No relevant evidence found in the document."
 
         except Exception as e:
-            logger.error(f"Error in LLM evidence retrieval: {str(e)}")
+            logger.error(f"Error in BERT evidence retrieval: {str(e)}")
             return f"Error retrieving evidence: {str(e)}"
 
 
-def retrieve_from_pdf_llm(pdf_file: str, query: str, k: int = 5) -> str:
+def retrieve_with_deberta(pdf_file: str, claim: str, top_k: int = 5) -> str:
     """
-    Wrapper function for LLM-based evidence retrieval
+    Wrapper function for DeBERTa-based evidence retrieval
     Compatible with the existing BM25S interface
 
     Args:
         pdf_file: Path to PDF file
-        query: Query/claim to find evidence for
-        k: Number of evidence passages to retrieve
+        claim: Claim to find evidence for
+        top_k: Number of evidence passages to retrieve
 
     Returns:
         Retrieved evidence text
     """
     # Initialize retriever (in production, this should be cached)
-    retriever = LLMEvidenceRetriever()
-    return retriever.retrieve_evidence(pdf_file, query, k)
+    retriever = BERTRetriever()
+    return retriever.retrieve_evidence(pdf_file, claim, top_k)
 
 
 # Alternative lightweight model for faster inference
-class LightweightLLMRetriever(LLMEvidenceRetriever):
+class DistilBERTRetriever(BERTRetriever):
     """
     Lightweight version using smaller, faster models
     """
@@ -221,17 +221,17 @@ class LightweightLLMRetriever(LLMEvidenceRetriever):
         super().__init__(model_name="distilbert-base-cased-distilled-squad")
 
 
-def retrieve_from_pdf_llm_fast(pdf_file: str, query: str, k: int = 5) -> str:
+def retrieve_with_distilbert(pdf_file: str, claim: str, top_k: int = 5) -> str:
     """
-    Fast LLM-based evidence retrieval using lightweight model
+    Fast DistilBERT-based evidence retrieval
 
     Args:
         pdf_file: Path to PDF file
-        query: Query/claim to find evidence for
-        k: Number of evidence passages to retrieve
+        claim: Claim to find evidence for
+        top_k: Number of evidence passages to retrieve
 
     Returns:
         Retrieved evidence text
     """
-    retriever = LightweightLLMRetriever()
-    return retriever.retrieve_evidence(pdf_file, query, k)
+    retriever = DistilBERTRetriever()
+    return retriever.retrieve_evidence(pdf_file, claim, top_k)

retrieval.py → retrieval_bm25s.py

@@ -5,7 +5,7 @@ from nltk.tokenize import sent_tokenize
 import bm25s
 
 
-def retrieve_from_pdf(pdf_file, query, k=10):
+def retrieve_with_bm25s(pdf_file, claim, top_k=10):
 
     # Get PDF file as binary
     with open(pdf_file, mode="rb") as f:
@@ -35,12 +35,12 @@ def retrieve_from_pdf(pdf_file, query, k=10):
     # Initialize the BM25 model
     retriever = bm25s.BM25()
     retriever.index(corpus_tokens, show_progress=False)
-    # Tokenize the query
-    query_tokens = bm25s.tokenize(query)
+    # Tokenize the claim
+    query_tokens = bm25s.tokenize(claim)
 
-    # Get top-k results
+    # Get top k results
     # Use int(k) in case we get str value (as in retrieval example)
-    results, scores = retriever.retrieve(query_tokens, corpus=corpus, k=int(k))
+    results, scores = retriever.retrieve(query_tokens, corpus=corpus, k=int(top_k))
     ## Print results
     # for i in range(results.shape[1]):
     #    doc, score = results[0, i], scores[0, i]

retrieval_gpt.py

@@ -0,0 +1,58 @@
+from openai import OpenAI
+import os
+from typing import Tuple
+
+
+def retrieve_with_gpt(pdf_file: str, claim: str) -> Tuple[str, int, int]:
+    """
+    Retrieve evidence from PDF using GPT
+
+    Args:
+        pdf_file: Path to PDF file
+        claim: Claim to find evidence for
+
+    Returns:
+        Tuple with retrieved evidence text, prompt tokens, and completion tokens
+    """
+
+    model = "gpt-4o-mini-2024-07-18"
+
+    prompt = """Retrieve sentences from the PDF (title, abstract, text, sections, not References/Bibliography) to support or refute this claim. \
+    Summarize any information from images. \
+    Respond only with verbatim sentences from the text and/or summarized sentences from images. \
+    If no conclusive evidence is found, respond with the five sentences that are most relevant to the claim. \
+    Combine all sentences into one response without quotation marks or line numbers. \
+    """
+
+    prompt = "".join([prompt, f"CLAIM: {claim}"])
+
+    client = OpenAI()
+
+    file = client.files.create(file=open(pdf_file, "rb"), purpose="user_data")
+
+    completion = client.chat.completions.create(
+        model=model,
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "file",
+                        "file": {
+                            "file_id": file.id,
+                        },
+                    },
+                    {
+                        "type": "text",
+                        "text": prompt,
+                    },
+                ],
+            }
+        ],
+    )
+
+    return (
+        completion.choices[0].message.content,
+        completion.usage.prompt_tokens,
+        completion.usage.completion_tokens,
+    )