Update app.py

app.py

@@ -1,69 +1,44 @@
-import zipfile
-import os
-import pandas as pd
-import torch
-from sentence_transformers import SentenceTransformer, util
-from transformers import pipeline
 import streamlit as st
-
-st.set_page_config(page_title="News Fact Checker", page_icon="📰")
-
-# Step 1: Unzip dataset
-@st.cache_data
-def extract_dataset():
-    zip_path = "climate.zip"
-    extract_dir = "climate_extracted"
-
-    if not os.path.exists(extract_dir):
-        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-            zip_ref.extractall(extract_dir)
-
-    train_path = os.path.join(extract_dir, "climate", "train")
-
-    # Try CSV or TSV format detection
-    try:
-        df = pd.read_csv(train_path, header=None)
-    except:
-        df = pd.read_csv(train_path, sep='\t', header=None)
-
-    df.columns = ["label", "title", "description"]
-    df["text"] = df["title"].astype(str) + ". " + df["description"].astype(str)
-    return df.head(1000)
-
-# Step 2: Load models
-@st.cache_resource
-def load_models():
-    embedder = SentenceTransformer('all-MiniLM-L6-v2')
-    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
-    return embedder, summarizer
-
-st.title("📰 News Fact Checker")
-st.markdown("Enter a **claim** about climate or news events. We'll pull relevant facts from real news and summarize them for you.")
-
-# Step 3: User input
-claim = st.text_input("🔍 Enter your claim:")
-data = extract_dataset()
-embedder, summarizer = load_models()
-
-# Step 4: Fact checking
-if claim:
-    with st.spinner("Searching news..."):
-        corpus = data["text"].tolist()
-        corpus_embeddings = embedder.encode(corpus, convert_to_tensor=True)
-        query_embedding = embedder.encode(claim, convert_to_tensor=True)
-
-        hits = util.semantic_search(query_embedding, corpus_embeddings, top_k=3)[0]
-        top_passages = [corpus[hit['corpus_id']] for hit in hits]
-
-        combined = " ".join(top_passages)
-        if len(combined) > 1024:
-            combined = combined[:1024]
-
-        summary = summarizer(combined, max_length=150, min_length=40, do_sample=False)[0]["summary_text"]
-
-    st.markdown("### ✅ Summary Based on News")
-    st.success(summary)
-
-    with st.expander("🔎 View Related News Snippets"):
-        for i, passage in enumerate(top_passages, 1):
-            st.markdown(f"**Snippet {i}:** {passage}")
+import pandas as pd
+from datasets import load_dataset
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+
+# Load AG News dataset from Hugging Face
+dataset = load_dataset("kk0105/ag-news", split="train")
+
+# Tokenizer and Model setup for RAG
+tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
+retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="default")
+model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
+
+# Function to generate response using RAG
+def generate_answer(query):
+    # Tokenize input query
+    inputs = tokenizer(query, return_tensors="pt")
+
+    # Retrieve relevant documents from dataset
+    input_ids = inputs["input_ids"]
+    question_embedding = retriever.compute_question_embeddings(input_ids)
+    context_input_ids = retriever.retrieve(input_ids, question_embedding)
+
+    # Generate an answer using the retrieved context
+    outputs = model.generate(input_ids=input_ids, context_input_ids=context_input_ids)
+    
+    # Decode the answer and return it
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return answer
+
+# Streamlit interface
+st.title("News Fact Checker")
+st.write("""
+    **Welcome to the News Fact Checker!**
+    Input a claim or question about a news topic, and we will verify or refute it based on recent news snippets.
+""")
+
+# User input for claim
+user_claim = st.text_input("Enter your claim or question:")
+
+if user_claim:
+    with st.spinner('Fetching relevant news snippets...'):
+        answer = generate_answer(user_claim)
+    st.write(f"**Fact Check Answer:** {answer}")