PS - 2: Hog RAGger

Dataset link: Hog RAGger Dataset
Dataset info:
Corpus.json: Contains collection of text of many articles scraped from the web. Contains following information about each article: title, author, source of article, publishing date, category, article url and article body (the text content).
Train.json: Contains output of a SOTA RAG pipeline. It contains 2556 queries, with evidence for each query distributed across 2 to 4 documents. The queries also involve document metadata, reflecting complex scenarios commonly found in real-world RAG applications. 

The goal is to develop a RAG pipeline that would answer the user’s question from the given corpus. The queries would require evidence from 2 to 4 documents and may also require information from document metadata in order to be answered. The output of your pipeline should be formatted the same as outputs in train.json and hence should contain the following:

query: The question posed to the pipeline, summarizing the topic or information being sought.
answer: The answer or inferred response to the query, generated based on the available evidence.
question_type: The type of question being asked, in this case, it's an "inference_query," indicating that the system inferred the answer from the context of the query and evidence.
evidence_list: A list of sources that provide supporting information for the answer. The list contains the following keys:
title: The title of the article or source that is referenced.
author: The name of the author or authors of the source.
url: The link to the source where the evidence is found.
source: The publication or website that published the source.
category: The general category or topic area of the source (e.g., technology, business).
published_at: The publication date and time of the source.
fact: A specific piece of information or fact extracted from the source that supports the query's answer.
Sample Input: "Who is the individual associated with the cryptocurrency industry facing a criminal trial on fraud and conspiracy charges, as reported by both The Verge and TechCrunch, and is accused by prosecutors of committing fraud for personal gain?"
Sample Output: 


Considerations:
Usage of any proprietary LLM/VLM (like Gemini, Claude, ChatGPT etc) or apis of such softwares is strictly forbidden. Use open source models instead.
Fine tuning is not allowed.
A small web app for demo would result in bonus points.

import streamlit as st
import json
from typing import List, Dict, Any
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load corpus and create embeddings
@st.cache_resource
def load_corpus(file_path: str) -> List[Dict[str, Any]]:
    with open(file_path, 'r') as f:
        return json.load(f)

@st.cache_data
def preprocess_corpus(corpus: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    processed_corpus = []
    for doc in corpus:
        processed_doc = {
            'title': doc['title'],
            'author': doc['author'],
            'url': doc['url'],
            'source': doc['source'],
            'category': doc['category'],
            'published_at': doc['published_at'],
            'content': doc['title'] + ' ' + doc['body']
        }
        processed_corpus.append(processed_doc)
    return processed_corpus

@st.cache_resource
def create_embeddings(corpus: List[Dict[str, Any]], model_name: str = 'sentence-transformers/all-MiniLM-L6-v2') -> np.ndarray:
    model = SentenceTransformer(model_name)
    return model.encode([doc['content'] for doc in corpus])

def retrieve_documents(query: str, corpus: List[Dict[str, Any]], embeddings: np.ndarray, model_name: str = 'sentence-transformers/all-MiniLM-L6-v2', top_k: int = 4) -> List[Dict[str, Any]]:
    model = SentenceTransformer(model_name)
    query_embedding = model.encode([query])
    similarities = cosine_similarity(query_embedding, embeddings)[0]
    top_indices = np.argsort(similarities)[-top_k:][::-1]
    return [corpus[i] for i in top_indices]

@st.cache_resource
def load_qa_model(model_name: str = 'google/flan-t5-base'):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    return tokenizer, model

def generate_answer(query: str, relevant_docs: List[Dict[str, Any]], tokenizer, model) -> str:
    context = " ".join([doc['content'] for doc in relevant_docs])
    input_text = f"Query: {query}\nContext: {context}\nAnswer:"
    
    inputs = tokenizer(input_text, return_tensors="pt", max_length=1024, truncation=True)
    outputs = model.generate(**inputs, max_length=150, num_return_sequences=1, temperature=0.7)
    
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def rag_pipeline(query: str, corpus: List[Dict[str, Any]], embeddings: np.ndarray, tokenizer, model) -> Dict[str, Any]:
    relevant_docs = retrieve_documents(query, corpus, embeddings)
    answer = generate_answer(query, relevant_docs, tokenizer, model)
    
    evidence_list = []
    for doc in relevant_docs:
        evidence = {
            'title': doc['title'],
            'author': doc['author'],
            'url': doc['url'],
            'source': doc['source'],
            'category': doc['category'],
            'published_at': doc['published_at'],
            'fact': doc['content'][:200] + '...'
        }
        evidence_list.append(evidence)
    
    return {
        'query': query,
        'answer': answer,
        'question_type': 'inference_query',
        'evidence_list': evidence_list
    }

# Streamlit app
st.set_page_config(page_title="RAG Query Processor", page_icon="🔍", layout="wide")

st.title("RAG Query Processor")

# Load data and models
corpus = load_corpus('corpus.json')
processed_corpus = preprocess_corpus(corpus)
embeddings = create_embeddings(processed_corpus)
tokenizer, model = load_qa_model()

# Query input
query = st.text_input("Enter your query:", "What is the impact of COVID-19 on the economy?")

if st.button("Process Query"):
    with st.spinner("Processing query..."):
        result = rag_pipeline(query, processed_corpus, embeddings, tokenizer, model)
    
    # Display results
    st.subheader("Answer")
    st.write(result['answer'])
    
    st.subheader("Evidence")
    for evidence in result['evidence_list']:
        with st.expander(evidence['title']):
            st.write(f"**Author:** {evidence['author']}")
            st.write(f"**Source:** {evidence['source']}")
            st.write(f"**Category:** {evidence['category']}")
            st.write(f"**Published:** {evidence['published_at']}")
            st.write(f"**Excerpt:** {evidence['fact']}")
            st.write(f"[Read More]({evidence['url']})")

st.sidebar.title("About")
st.sidebar.info("This app uses a RAG (Retrieval-Augmented Generation) pipeline to answer queries based on a corpus of documents.")
Generate a readme.md as the answer to the problem statement describing what we have done and how to run everytihng in detaail