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
)
tokenizer = AutoTokenizer.from_pretrained('google/flan-t5-base')
classifier = pipeline(
    "text-classification",
    model=model,
    tokenizer=tokenizer,
    trust_remote_code=True
)

# 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

    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 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 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.")