app.py

import streamlit as st
import os
import requests
from langdetect import detect
from PyPDF2 import PdfReader
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.neighbors import NearestNeighbors
import numpy as np

# Load the Hugging Face token from environment variables
huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")  # Replace with your Hugging Face token

# Function to query the Hugging Face API
def query_huggingface_model(prompt, max_new_tokens=1000, temperature=0.7, top_k=50):
    model_name = "HuggingFaceH4/zephyr-7b-alpha"  # Replace with your preferred model
    api_url = f"https://api-inference.huggingface.co/models/{model_name}"
    headers = {"Authorization": f"Bearer {huggingface_token}"}
    payload = {
        "inputs": prompt,
        "parameters": {
            "max_new_tokens": max_new_tokens,
            "temperature": temperature,
            "top_k": top_k,
        },
    }
    response = requests.post(api_url, headers=headers, json=payload)
    if response.status_code == 200:
        return response.json()[0]["generated_text"]
    else:
        st.error(f"Error: {response.status_code} - {response.text}")
        return None

# Function to detect language
def detect_language(text):
    try:
        return detect(text)
    except:
        return "en"  # Default to English if detection fails

# Function to extract text from PDF with line and page numbers
def extract_text_from_pdf(pdf_file):
    pdf_reader = PdfReader(pdf_file)
    text_data = []
    for page_num, page in enumerate(pdf_reader.pages):
        lines = page.extract_text().split('\n')
        for line_num, line in enumerate(lines):
            text_data.append({
                "page": page_num + 1,
                "line": line_num + 1,
                "content": line
            })
    return text_data

# Function to search for query in PDF content
def search_pdf_content(pdf_text_data, query):
    results = []
    for entry in pdf_text_data:
        if query.lower() in entry["content"].lower():
            results.append(entry)
    return results

# Function to split text into chunks
def split_text_into_chunks(text, chunk_size=500):
    words = text.split()
    chunks = [" ".join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
    return chunks

# Function to compute cosine similarity between query and document chunks
def compute_cosine_similarity(query, chunks):
    vectorizer = TfidfVectorizer()
    tfidf_matrix = vectorizer.fit_transform([query] + chunks)
    cosine_similarities = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[1:]).flatten()
    return cosine_similarities

# Function to find KNN-based similar documents
def find_knn_similar_documents(query, chunks, k=5):
    vectorizer = TfidfVectorizer()
    tfidf_matrix = vectorizer.fit_transform([query] + chunks)
    knn = NearestNeighbors(n_neighbors=k, metric="cosine")
    knn.fit(tfidf_matrix[1:])
    distances, indices = knn.kneighbors(tfidf_matrix[0:1])
    return indices.flatten(), distances.flatten()

# Default system prompts for each query translation method
DEFAULT_SYSTEM_PROMPTS = {
    "Multi-Query": """You are an AI language model assistant. Your task is to generate five 
different versions of the given user question to retrieve relevant documents from a vector 
database. By generating multiple perspectives on the user question, your goal is to help
the user overcome some of the limitations of the distance-based similarity search. 
Provide these alternative questions separated by newlines. Original question: {question}""",
    "RAG Fusion": """You are an AI language model assistant. Your task is to combine multiple 
queries into a single, refined query to improve retrieval accuracy. Original question: {question}""",
    "Decomposition": """You are an AI language model assistant. Your task is to break down 
the given user question into simpler sub-questions. Provide these sub-questions separated 
by newlines. Original question: {question}""",
    "Step Back": """You are an AI language model assistant. Your task is to refine the given 
user question by taking a step back and asking a more general question. Original question: {question}""",
    "HyDE": """You are an AI language model assistant. Your task is to generate a hypothetical 
document that would be relevant to the given user question. Original question: {question}""",
}

# Streamlit App
def main():
    st.title("RAG Model with Advanced Query Translation and Indexing")
    st.write("Enter a prompt and get a response from the model.")

    # Sidebar for options
    st.sidebar.title("Options")

    # PDF Upload
    st.sidebar.header("Upload PDF")
    pdf_file = st.sidebar.file_uploader("Upload a PDF file", type="pdf")

    # Query Translation Options
    st.sidebar.header("Query Translation")
    query_translation = st.sidebar.selectbox(
        "Select Query Translation Method",
        ["Multi-Query", "RAG Fusion", "Decomposition", "Step Back", "HyDE"]
    )

    # Indexing Options
    st.sidebar.header("Indexing")
    indexing_method = st.sidebar.selectbox(
        "Select Indexing Method",
        ["Multi-Representation", "Raptors", "ColBERT"]
    )

    # Similarity Search Options
    st.sidebar.header("Similarity Search")
    similarity_method = st.sidebar.selectbox(
        "Select Similarity Search Method",
        ["Cosine Similarity", "KNN"]
    )
    if similarity_method == "KNN":
        k_value = st.sidebar.slider("Select K Value", 1, 10, 5)

    # LLM Parameters
    st.sidebar.header("LLM Parameters")
    max_new_tokens = st.sidebar.slider("Max New Tokens", 10, 1000, 1000)
    temperature = st.sidebar.slider("Temperature", 0.1, 1.0, 0.7)
    top_k = st.sidebar.slider("Top K", 1, 100, 50)

    # System Prompt
    st.sidebar.header("System Prompt")
    default_system_prompt = DEFAULT_SYSTEM_PROMPTS[query_translation]
    system_prompt = st.sidebar.text_area("System Prompt", default_system_prompt)

    # Main Content
    st.header("Input Prompt")
    prompt = st.text_input("Enter your prompt:")
    if prompt:
        st.write("**Prompt:**", prompt)

        # Detect Language
        language = detect_language(prompt)
        st.write(f"**Detected Language:** {language}")

        # Query Translation
        if st.button("Apply Query Translation"):
            st.write(f"**Applied Query Translation Method:** {query_translation}")
            # Format the system prompt with the user's question
            formatted_prompt = system_prompt.format(question=prompt)
            st.write("**Formatted System Prompt:**", formatted_prompt)

            # Query the Hugging Face model for query translation
            translated_queries = query_huggingface_model(formatted_prompt, max_new_tokens, temperature, top_k)
            if translated_queries:
                st.write("**Translated Queries:**")
                st.write(translated_queries.split("\n")[-1])  # Print only the updated question part

        # Indexing
        if st.button("Apply Indexing"):
            st.write(f"**Applied Indexing Method:** {indexing_method}")
            if pdf_file is not None:
                # Extract and search PDF content
                pdf_text_data = extract_text_from_pdf(pdf_file)
                search_results = search_pdf_content(pdf_text_data, prompt)

                if search_results:
                    st.write("**Relevant Content from PDF:**")
                    for result in search_results:
                        st.write(f"**Page {result['page']}, Line {result['line']}:** {result['content']}")

                    # Split text into chunks
                    chunks = split_text_into_chunks("\n".join([result["content"] for result in search_results]))
                    st.write("**Chunks Obtained from PDF:**")
                    for i, chunk in enumerate(chunks):
                        st.write(f"**Chunk {i + 1}:** {chunk}")

                    # Perform similarity search
                    if similarity_method == "Cosine Similarity":
                        st.write("**Cosine Similarity Results:**")
                        cosine_similarities = compute_cosine_similarity(prompt, chunks)
                        for i, similarity in enumerate(cosine_similarities):
                            st.write(f"**Chunk {i + 1} Similarity:** {similarity:.4f}")
                    elif similarity_method == "KNN":
                        st.write(f"**KNN Results (k={k_value}):**")
                        indices, distances = find_knn_similar_documents(prompt, chunks, k_value)
                        for i, (index, distance) in enumerate(zip(indices, distances)):
                            st.write(f"**Chunk {index + 1} Distance:** {distance:.4f}")
                else:
                    st.write("**No relevant content found in the PDF.**")
            else:
                st.write("**No PDF uploaded.**")

        # Generate Response
        if st.button("Generate Response"):
            if pdf_file is not None:
                # Extract and search PDF content
                pdf_text_data = extract_text_from_pdf(pdf_file)
                search_results = search_pdf_content(pdf_text_data, prompt)

                if search_results:
                    st.write("**Relevant Content from PDF:**")
                    for result in search_results:
                        st.write(f"**Page {result['page']}, Line {result['line']}:** \"{result['content']}\"")

                    # Generate response based on PDF content
                    pdf_context = "\n".join([result["content"] for result in search_results])
                    response = query_huggingface_model(f"Based on the following context:\n{pdf_context}\n\nAnswer this question: {prompt}", max_new_tokens, temperature, top_k)
                else:
                    st.write("**No relevant content found in the PDF. Generating response without PDF context.**")
                    response = query_huggingface_model(prompt, max_new_tokens, temperature, top_k)
            else:
                st.write("**No PDF uploaded. Generating response without PDF context.**")
                response = query_huggingface_model(prompt, max_new_tokens, temperature, top_k)

            if response:
                st.write("**Response:**", response)

if __name__ == "__main__":
    main()