add strict relevancy and scite badges and reranking

README.md

@@ -9,5 +9,3 @@ app_file: app.py
 pinned: false
 license: cc-by-2.0
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -2,15 +2,38 @@ import streamlit as st
 from transformers import pipeline
 import requests
 from bs4 import BeautifulSoup
+from nltk.corpus import stopwords
+import nltk
+import string
+from streamlit.components.v1 import html
+from sentence_transformers.cross_encoder import CrossEncoder as CE
+import numpy as np
+from typing import List, Tuple
+import torch
+
+class CrossEncoder:
+    def __init__(self, model_path: str, **kwargs):
+        self.model = CE(model_path, **kwargs)
+
+    def predict(self, sentences: List[Tuple[str,str]], batch_size: int = 32, show_progress_bar: bool = True) -> List[float]:
+        return self.model.predict(
+            sentences=sentences,
+            batch_size=batch_size,
+            show_progress_bar=show_progress_bar)
+
 
 SCITE_API_KEY = st.secrets["SCITE_API_KEY"]
 
+
 def remove_html(x):
     soup = BeautifulSoup(x, 'html.parser')
     text = soup.get_text()
     return text
 
-def search(term, limit=25):
+
+def search(term, limit=10, clean=True, strict=True):
+    term = clean_query(term, clean=clean, strict=strict)
+    # heuristic, 2 searches strict and not? and then merge?
     search = f"https://api.scite.ai/search?mode=citations&term={term}&limit={limit}&offset=0&user_slug=domenic-rosati-keW5&compute_aggregations=false"
     req = requests.get(
         search,
@@ -19,8 +42,9 @@ def search(term, limit=25):
         }
     )
     return (
-        remove_html('\n'.join(['\n'.join([cite['snippet'] for cite in doc['citations']]) for doc in req.json()['hits']])),
-        [(doc['doi'], doc['citations'], doc['title']) for doc in req.json()['hits']]
+        [remove_html('\n'.join([cite['snippet'] for cite in doc['citations']])) for doc in req.json()['hits']],
+        [(doc['doi'], doc['citations'], doc['title'])
+         for doc in req.json()['hits']]
     )
 
 
@@ -39,25 +63,37 @@ def find_source(text, docs):
                 'source_title': doc[2],
                 'source_link': f"https://scite.ai/reports/{doc[0]}"
             }
-    return {
-            'citation_statement': '',
-            'text': text,
-            'from': '',
-            'supporting': '',
-            'source_title': '',
-            'source_link': ''
-        }
+    return None
+
 
 @st.experimental_singleton
 def init_models():
-    question_answerer = pipeline("question-answering", model='sultan/BioM-ELECTRA-Large-SQuAD2-BioASQ8B')
-    return question_answerer
+    nltk.download('stopwords')
+    stop = set(stopwords.words('english') + list(string.punctuation))
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    question_answerer = pipeline(
+        "question-answering", model='sultan/BioM-ELECTRA-Large-SQuAD2-BioASQ8B',
+        device=device
+    )
+    reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device=device)
+    return question_answerer, reranker, stop, device
+
+qa_model, reranker, stop, device = init_models()
 
-qa_model = init_models()
+def clean_query(query, strict=True, clean=True):
+    operator = ' '
+    if strict:
+        operator = ' AND '
+    query = operator.join(
+        [i for i in query.lower().split(' ') if clean and i not in stop])
+    if clean:
+        query = query.translate(str.maketrans('', '', string.punctuation))
+    return query
 
 
-def card(title, context, score, link):
-    return st.markdown(f"""
+
+def card(title, context, score, link, supporting):
+    st.markdown(f"""
     <div class="container-fluid">
         <div class="row align-items-start">
              <div  class="col-md-12 col-sm-12">
@@ -72,6 +108,22 @@ def card(title, context, score, link):
         </div>
      </div>
         """, unsafe_allow_html=True)
+    html(f"""
+    <div
+                    class="scite-badge"
+                    data-doi="{supporting}"
+                    data-layout="horizontal"
+                    data-show-zero="false"
+                    data-show-labels="false"
+                    data-tally-show="true"
+                />
+    <script
+    async
+    type="application/javascript"
+    src="https://cdn.scite.ai/badge/scite-badge-latest.min.js">
+  </script>
+  """, width=None, height=42, scrolling=False)
+
 
 st.title("Scientific Question Answering with Citations")
 
@@ -85,8 +137,14 @@ st.markdown("""
 """, unsafe_allow_html=True)
 
 def run_query(query):
-    context, orig_docs = search(query)
-    if not context.strip():
+    if device == 'cpu':
+        limit = 50
+        context_limit = 10
+    else:
+        limit = 100
+        context_limit = 25
+    contexts, orig_docs = search(query, limit=limit)
+    if len(contexts) == 0 or not ''.join(contexts).strip():
         return st.markdown("""
         <div class="container-fluid">
         <div class="row align-items-start">
@@ -97,35 +155,44 @@ def run_query(query):
      </div>
         """, unsafe_allow_html=True)
 
+    sentence_pairs = [[query, context] for context in contexts]
+    scores = reranker.predict(sentence_pairs, batch_size=limit, show_progress_bar=False)
+    hits = {contexts[idx]: scores[idx] for idx in range(len(scores))}
+    sorted_contexts = [k for k,v in sorted(hits.items(), key=lambda x: x[0], reverse=True)]
+
+    context = '\n'.join(sorted_contexts[:context_limit])
     results = []
     model_results = qa_model(question=query, context=context, top_k=10)
     for result in model_results:
         support = find_source(result['answer'], orig_docs)
+        if not support:
+            continue
         results.append({
             "answer": support['text'],
             "title": support['source_title'],
             "link": support['source_link'],
             "context": support['citation_statement'],
-            "score": result['score']
+            "score": result['score'],
+            "doi": support["supporting"]
         })
 
-
-
     sorted_result = sorted(results, key=lambda x: x['score'], reverse=True)
     sorted_result = list({
         result['context']: result for result in sorted_result
     }.values())
-    sorted_result = sorted(sorted_result, key=lambda x: x['score'], reverse=True)
-
+    sorted_result = sorted(
+        sorted_result, key=lambda x: x['score'], reverse=True)
 
     for r in sorted_result:
         answer = r["answer"]
-        ctx = remove_html(r["context"]).replace(answer, f"<mark>{answer}</mark>").replace('<cite', '<a').replace('</cite', '</a').replace('data-doi="', 'href="https://scite.ai/reports/')
-        title = r["title"].replace("_", " ")
+        ctx = remove_html(r["context"]).replace(answer, f"<mark>{answer}</mark>").replace(
+            '<cite', '<a').replace('</cite', '</a').replace('data-doi="', 'href="https://scite.ai/reports/')
+        title = r.get("title", '').replace("_", " ")
         score = round(r["score"], 4)
-        card(title, ctx, score, r['link'])
+        card(title, ctx, score, r['link'], r['doi'])
 
 query = st.text_input("Ask scientific literature a question", "")
 
 if query != "":
-    run_query(query)
+    with st.spinner('Loading...'):
+        run_query(query)

requirements.txt

@@ -3,3 +3,6 @@ requests
 beautifulsoup4
 streamlit==1.2.0
 torch
+nltk
+sentence_transformers
+numpy