add treshold for predictions

app.py

@@ -11,6 +11,8 @@ import numpy as np
 from typing import List, Tuple
 import torch
 
+SCITE_API_KEY = st.secrets["SCITE_API_KEY"]
+
 class CrossEncoder:
     def __init__(self, model_path: str, **kwargs):
         self.model = CE(model_path, **kwargs)
@@ -22,18 +24,21 @@ class CrossEncoder:
             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
 
 
+# 4 searches: strict y/n, supported y/n
+# deduplicate
+# search per query
+
+
 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?
+    # https://api.scite.ai/search?mode=citations&term=unit%20testing%20software&limit=10&date_from=2000&date_to=2022&offset=0&supporting_from=1&contrasting_from=0&contrasting_to=0&user_slug=domenic-rosati-keW5&compute_aggregations=true
     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,
@@ -67,6 +72,7 @@ def find_source(text, docs):
                 'source_title': doc[2],
                 'source_link': f"https://scite.ai/reports/{doc[0]}"
             }
+    print("None found for", text)
     return None
 
 
@@ -79,7 +85,7 @@ def init_models():
         "question-answering", model='sultan/BioM-ELECTRA-Large-SQuAD2-BioASQ8B',
         device=device
     )
-    reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2', device=device)
+    reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device=device)
     queryexp_tokenizer = AutoTokenizer.from_pretrained("doc2query/all-with_prefix-t5-base-v1")
     queryexp_model = AutoModelWithLMHead.from_pretrained("doc2query/all-with_prefix-t5-base-v1")
     return question_answerer, reranker, stop, device, queryexp_model, queryexp_tokenizer
@@ -98,7 +104,6 @@ def clean_query(query, strict=True, clean=True):
     return query
 
 
-
 def card(title, context, score, link, supporting):
     st.markdown(f"""
     <div class="container-fluid">
@@ -138,7 +143,7 @@ st.write("""
 Ask a scientific question and get an answer drawn from [scite.ai](https://scite.ai) corpus of over 1.1bn citation statements.
 Answers are linked to source documents containing citations where users can explore further evidence from scientific literature for the answer.
 
-For example try: Are tanning beds safe to use? Does size of venture capital fund correlate with returns?
+For example try: Do tanning beds cause cancer?
 """)
 
 st.markdown("""
@@ -146,26 +151,27 @@ st.markdown("""
 """, unsafe_allow_html=True)
 
 with st.expander("Settings (strictness, context limit, top hits)"):
+    confidence_threshold = st.slider('Confidence threshold for answering questions? This number represents how confident the model should be in the answers it gives. The number is out of 100%', 0, 100, 1)
     strict_mode = st.radio(
         "Query mode? Strict means all words must match in source snippet. Lenient means only some words must match.",
         ('lenient', 'strict'))
     use_reranking = st.radio(
         "Use Reranking? Reranking will rerank the top hits using semantic similarity of document and query.",
         ('yes', 'no'))
+    top_hits_limit = st.slider('Top hits? How many documents to use for reranking. Larger is slower but higher quality', 10, 300, 200 if torch.cuda.is_available() else 50)
+    context_lim = st.slider('Context limit? How many documents to use for answering from. Larger is slower but higher quality', 10, 300, 25 if torch.cuda.is_available() else 10)
     use_query_exp = st.radio(
         "(Experimental) use query expansion? Right now it just recommends queries",
         ('yes', 'no'))
-    top_hits_limit = st.slider('Top hits? How many documents to use for reranking. Larger is slower but higher quality', 10, 300, 200 if torch.cuda.is_available() else 100)
-    context_lim = st.slider('Context limit? How many documents to use for answering from. Larger is slower but higher quality', 10, 300, 25 if torch.cuda.is_available() else 10)
+    suggested_queries = st.slider('Number of suggested queries to use', 0, 10, 5)
 
 def paraphrase(text, max_length=128):
     input_ids = queryexp_tokenizer.encode(text, return_tensors="pt", add_special_tokens=True)
-    generated_ids = queryexp_model.generate(input_ids=input_ids, num_return_sequences=5, num_beams=5, max_length=max_length)
+    generated_ids = queryexp_model.generate(input_ids=input_ids, num_return_sequences=suggested_queries or 5, num_beams=suggested_queries or 5, max_length=max_length)
     queries = set([queryexp_tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in generated_ids])
     preds = '\n * '.join(queries)
     return preds
 
-
 def run_query(query):
     if use_query_exp == 'yes':
         query_exp = paraphrase(f"question2question: {query}")
@@ -186,7 +192,6 @@ If you are not getting good results try one of:
         </div>
      </div>
         """, unsafe_allow_html=True)
-
     if use_reranking == 'yes':
         sentence_pairs = [[query, context] for context in contexts]
         scores = reranker.predict(sentence_pairs, batch_size=limit, show_progress_bar=False)
@@ -195,7 +200,6 @@ If you are not getting good results try one of:
         context = '\n'.join(sorted_contexts[:context_limit])
     else:
         context = '\n'.join(contexts[:context_limit])
-
     results = []
     model_results = qa_model(question=query, context=context, top_k=10)
     for result in model_results:
@@ -210,14 +214,23 @@ If you are not getting good results try one of:
             "score": result['score'],
             "doi": support["supporting"]
         })
-
-    sorted_result = sorted(results, key=lambda x: x['score'], reverse=True)
+    sorted_result = sorted(results, key=lambda x: x['score'])
     sorted_result = list({
         result['context']: result for result in sorted_result
     }.values())
     sorted_result = sorted(
         sorted_result, key=lambda x: x['score'], reverse=True)
 
+    if confidence_threshold == 0:
+        threshold = 0
+    else:
+        threshold = (confidence_threshold or 10) / 100
+
+    sorted_result = filter(
+        lambda x: x['score'] > threshold,
+        sorted_result
+    )
+
     for r in sorted_result:
         answer = r["answer"]
         ctx = remove_html(r["context"]).replace(answer, f"<mark>{answer}</mark>").replace(
@@ -227,7 +240,6 @@ If you are not getting good results try one of:
         card(title, ctx, score, r['link'], r['doi'])
 
 query = st.text_input("Ask scientific literature a question", "")
-
 if query != "":
     with st.spinner('Loading...'):
         run_query(query)