Update app.py

app.py

@@ -1,8 +1,11 @@
 import streamlit as st
 from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
+import wikipediaapi
 
 # Load the hallucination evaluation model
-model = AutoModelForSequenceClassification.from_pretrained("vectara/hallucination_evaluation_model", trust_remote_code=True)
+model = AutoModelForSequenceClassification.from_pretrained(
+    "vectara/hallucination_evaluation_model", trust_remote_code=True
+)
 tokenizer = AutoTokenizer.from_pretrained('google/flan-t5-base')
 classifier = pipeline(
     "text-classification",
@@ -11,40 +14,39 @@ classifier = pipeline(
     trust_remote_code=True
 )
 
-# Function to generate a premise (fact-based reference)
-def generate_premise(query):
-    # This can use GPT, external APIs, or predefined logic
-    # Example using GPT or any fact-retrieval API:
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-
-    gpt_model = "google/flan-t5-base"  # Replace with your preferred model
-    gpt_tokenizer = AutoTokenizer.from_pretrained(gpt_model)
-    gpt = AutoModelForCausalLM.from_pretrained(gpt_model)
+# Function to generate a premise using Wikipedia
+def generate_premise_wikipedia(query):
+    """
+    Retrieves the most relevant Wikipedia summary as a premise for the given query.
+    """
+    wiki_wiki = wikipediaapi.Wikipedia('en')  # English Wikipedia
+    search_results = wiki_wiki.page(query)  # Search Wikipedia for the query
 
-    # Generate a fact-based premise
-    prompt = f"Generate a fact-based premise for the query: {query}"
-    inputs = gpt_tokenizer(prompt, return_tensors="pt")
-    outputs = gpt.generate(**inputs)
-    premise = gpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return premise
+    if search_results.exists():
+        # If the page exists, return the summary as the premise
+        return search_results.summary
+    else:
+        return "No relevant information found on Wikipedia."
 
 # Streamlit UI
 st.title("Hallucination Detection App")
-st.write("Enter a query, and the app will tell you if it is hallucinated or accurate based on factual information.")
+st.write("Enter a query, and the app will check if it is hallucinated or factual using Wikipedia as a factual source.")
 
 # Input query from user
 query = st.text_input("Enter your query:")
 
 if query:
-    # Generate a premise for the query
-    premise = generate_premise(query)
-
-    # Prepare input for hallucination evaluation
-    prompt = f"<pad> Determine if the hypothesis is true given the premise?\n\nPremise: {premise}\n\nHypothesis: {query}"
-    result = classifier([prompt], top_k=None)[0][0]
-
-    # Output results
-    st.write(f"**Premise:** {premise}")
-    st.write(f"**Hypothesis:** {query}")
-    st.write(f"**Result:** {result['label']} (Confidence: {result['score']:.2f})")
-
+    # Generate a premise for the query using Wikipedia
+    premise = generate_premise_wikipedia(query)
+
+    if premise != "No relevant information found on Wikipedia.":
+        # Prepare input for hallucination evaluation
+        prompt = f"<pad> Determine if the hypothesis is true given the premise?\n\nPremise: {premise}\n\nHypothesis: {query}"
+        result = classifier([prompt], top_k=None)[0][0]
+
+        # Output results
+        st.write(f"**Premise (from Wikipedia):** {premise}")
+        st.write(f"**Hypothesis:** {query}")
+        st.write(f"**Result:** {result['label']} (Confidence: {result['score']:.2f})")
+    else:
+        st.write("**Premise:** No relevant information found on Wikipedia.")