support for different NLI models; early stopping in entailment

Rock_fact_checker.py

@@ -76,6 +76,7 @@ def main():
         with st.spinner("🧠    Performing neural search on documents..."):
             try:
                 st.session_state.results = query(statement, RETRIEVER_TOP_K)
+                print(query)
                 time_end = time.time()
                 print(time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()))
                 print(f"elapsed time: {time_end - time_start}")

app_utils/backend_utils.py

@@ -77,9 +77,9 @@ def query(statement: str, retriever_top_k: int = 5):
         agg_neu += neu * doc.score
         agg_ent += ent * doc.score
 
-        # if in the first 3 documents there is a strong evidence of entailment/contradiction,
-        # there is non need to consider less relevant documents
-        if i == 2 and max(agg_con, agg_ent) / scores > 0.5:
+        # if in the first documents there is a strong evidence of entailment/contradiction,
+        # there is no need to consider less relevant documents
+        if max(agg_con, agg_ent) / scores > 0.5:
             results["documents"] = results["documents"][: i + 1]
             break
 

app_utils/config.py

@@ -1,7 +1,16 @@
+from curses.ascii import NL
+from logging import exception
+import streamlit as st
+
 INDEX_DIR = "data/index"
 STATEMENTS_PATH = "data/statements.txt"
 
 RETRIEVER_MODEL = "sentence-transformers/msmarco-distilbert-base-tas-b"
 RETRIEVER_MODEL_FORMAT = "sentence_transformers"
 RETRIEVER_TOP_K = 5
-NLI_MODEL = "valhalla/distilbart-mnli-12-1"
+
+try:
+    NLI_MODEL = st.secrets['NLI_MODEL']
+except:
+    NLI_MODEL = "valhalla/distilbart-mnli-12-1"
+print(f'Used NLI model: {NLI_MODEL}')

data/statements.txt

@@ -18,7 +18,6 @@ The White Stripes were a trio
 The White Stripes were composed by Jack White and Meg White
 Scorpions is a German trap band
 Sepultura is a heavy metal band
-Toxicity is a song by System of a down
 System of a down is a Italian band
 The Cure is a pop band
 Mick Jagger loves pasta
@@ -36,3 +35,14 @@ Guns N' Roses formed in 1987
 Izzy Stradlin took part in Guns N' Roses
 Grunge was born in Seattle
 Dave Grohl played in Nirvana and Foo Fighters
+John Michael Stipe was part of the R.E.M.
+Eric Clapton did some gigs with Dire Straits
+Mark Knopfler plays drums
+"La Bamba" is a hit by Ritchie Valens
+Nico took part in the Velvet Underground
+David Bowie appears in Twin Peaks
+The Cure made dark songs
+Cannibal Corpse is a pop punk band
+Slipknot wear masks
+Toto have sold many records
+The verve were a prominent English band