app.py

import gradio as gr
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from sentence_transformers import SentenceTransformer, util
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import re
import traceback
import torch
import os
from sentence_transformers import SentenceTransformer, util
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import re
import pandas as pd
import json


# Preprocessing text by lowercasing, removing punctuation, and extra spaces
def optimized_preprocess_text(text):
    text = text.lower()
    text = re.sub(r'[^\w\s]', '', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text

# Compute cosine similarity between two texts using TF-IDF
def optimized_compute_text_similarity(text1, text2):
    tfidf = TfidfVectorizer(stop_words='english', ngram_range=(1, 1))
    tfidf_matrix = tfidf.fit_transform([text1, text2])
    cosine_sim = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[1:2]).flatten()
    return cosine_sim[0]

# Compute SBERT similarity between question and context
def compute_sbert_similarity(question, context, model):
    embeddings = model.encode([question, context], convert_to_tensor=True)
    similarity = util.pytorch_cos_sim(embeddings[0], embeddings[1]).item()
    return similarity

# Use hybrid approach: TF-IDF to narrow down top N contexts, then SBERT for refined similarity
def hybrid_sbert_approach(question, filtered_contexts, model, top_n=10):
    tfidf = TfidfVectorizer(stop_words='english')
    contexts_combined = [question] + filtered_contexts
    tfidf_matrix = tfidf.fit_transform(contexts_combined)
    
    # Calculate TF-IDF similarity and rank contexts
    similarity_scores = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[1:]).flatten()
    ranked_contexts = [filtered_contexts[i] for i in similarity_scores.argsort()[::-1][:top_n]]
    
    # Refine using SBERT
    sbert_similarities = [compute_sbert_similarity(question, context, model) for context in ranked_contexts]
    ranked_by_sbert = sorted(zip(ranked_contexts, sbert_similarities), key=lambda x: x[1], reverse=True)
    
    return [context for context, _ in ranked_by_sbert]

# RAG with optimized SBERT function
def optimized_generate_rag_context(question, filtered_contexts, selected_context_window=2):
    hybrid_retrieved_contexts = hybrid_sbert_approach(question, filtered_contexts, sbert_model, top_n=int(selected_context_window))
    rag_context = "\n".join(hybrid_retrieved_contexts[:selected_context_window])
    return rag_context

# Extract unique contexts and filter them by length
def extract_and_filter_contexts(data, min_length=151, max_length=3706):
    unique_contexts = data['context'].unique()
    filtered_contexts = [context for context in unique_contexts if min_length <= len(context) <= max_length]
    return filtered_contexts

# Compute the TF-IDF matrix for the question and contexts
def compute_tfidf_and_similarity_scores(question, contexts):
    tfidf = TfidfVectorizer(stop_words='english')
    contexts_combined = [question] + contexts
    tfidf_matrix = tfidf.fit_transform(contexts_combined)
    
    # Calculate the cosine similarity scores
    similarity_scores = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[1:]).flatten()
    return tfidf_matrix, similarity_scores

# Rank contexts based on similarity scores
def rank_contexts_by_similarity(contexts, similarity_scores):
    ranked_indices = similarity_scores.argsort()[::-1]
    ranked_contexts = [contexts[i] for i in ranked_indices]
    ranked_scores = similarity_scores[ranked_indices]
    return ranked_contexts, ranked_scores

# Select the top contexts based on the selected window
def select_top_contexts(selected_context_window, ranked_contexts, ranked_scores):
    count = int(selected_context_window)
    top_contexts = ranked_contexts[:count]
    top_scores = ranked_scores[:count]
    return top_contexts, top_scores


# Helper function to maintain chat history and generate the response
def maintain_chat_history(message, chat_history):
    if chat_history is None:
        chat_history = []
    chat_history.append({"role": "user", "content": message})
    return chat_history

def generate_rag_context(question, filtered_contexts, selected_context_window = 3):
    tfidf_matrix, similarity_scores = compute_tfidf_and_similarity_scores(question, filtered_contexts)
    ranked_contexts, ranked_scores = rank_contexts_by_similarity(filtered_contexts, similarity_scores)
    top_contexts, top_scores = select_top_contexts(str(selected_context_window), ranked_contexts, ranked_scores)
    rag_context = "\n".join(top_contexts)
    return rag_context

def load_squad_data(filepath):
    with open(filepath, 'r') as f:
        squad_data = json.load(f)
    return squad_data



# Preprocess the data: extract contexts, questions, and answers from the SQuAD data
def raw_preprocess_data(squad_data):
    contexts = []
    questions = []
    answers = []

    for group in squad_data['data']:
        for passage in group['paragraphs']:
            context = passage['context']
            for qa in passage['qas']:
                question = qa['question']
                for answer in qa['answers']:
                    contexts.append(context)
                    questions.append(question)
                    # Make a copy to avoid modifying the original answer
                    answers.append({
                        'text': answer['text'],
                        'answer_start': answer['answer_start']
                    })

    return contexts, questions, answers


# Add the end index of the answer in the context
def add_end_idx(answers, contexts):
    for answer, context in zip(answers, contexts):
        gold_text = answer['text']
        start_idx = answer['answer_start']
        end_idx = start_idx + len(gold_text)

        if context[start_idx:end_idx] == gold_text:
            answer['answer_end'] = end_idx
        else:
            # Try to find the correct position if there's a mismatch
            for n in range(1, 30):
                if context[start_idx - n:end_idx - n] == gold_text:
                    answer['answer_start'] = start_idx - n
                    answer['answer_end'] = end_idx - n
                    break
                elif context[start_idx + n:end_idx + n] == gold_text:
                    answer['answer_start'] = start_idx + n
                    answer['answer_end'] = end_idx + n
                    break
            else:
                answer['answer_start'] = -1
                answer['answer_end'] = -1


# Create a DataFrame from the contexts, questions, and answers
def create_dataframe(contexts, questions, answers):
    data = pd.DataFrame({
        'context': contexts,
        'question': questions,
        'answer_text': [answer['text'] for answer in answers],
        'answer_start': [answer['answer_start'] for answer in answers],
        'answer_end': [answer.get('answer_end', -1) for answer in answers]
    })

    # Remove samples with -1 start index
    data = data[data['answer_start'] != -1].reset_index(drop=True)
    return data

# Check if a GPU (CUDA) is available; otherwise, use the CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# Loading the pre-trained SBERT model globally for efficiency
sbert_model = SentenceTransformer('all-MiniLM-L6-v2')

# Available models
electra_models = [
    "./models/fine_tuned_electra_model_1000",
    "./models/fine_tuned_electra_model_20000",
    "./models/fine_tuned_electra_model_5000",
    "./models/fine_tuned_electra_model_all"
]
other_models = [
    "./models/fine_tuned_bert_base_cased_1000",
    "./models/fine_tuned_bert_base_cased_all",
    "./models/fine_tuned_distilbert_base_uncased_10000",
    "./models/fine_tuned_distilgpt2_10000",
    "./models/fine_tuned_retro-reader_intensive_1000",
    "./models/fine_tuned_retro-reader_intensive_5000",
    "./models/fine_tuned_retro-reader_sketchy_1000"
]

DATA_DIR = './data'

# Load and preprocess data
squad_data = load_squad_data(DATA_DIR+ '/train-v1.1.json')
contexts, questions, answers = raw_preprocess_data(squad_data)
add_end_idx(answers, contexts)
data = create_dataframe(contexts, questions, answers)

# Function to generate a response with logging and custom content
def generate_response(message, chat_history, model_name, debug, rag, selected_context_window):
    try:
        if chat_history is None:
            chat_history = []
        context = message

        # Determine if the model is for question answering based on its name
        is_question_answering = "electra_model" in model_name

        # Initialize the tokenizer and model
        if is_question_answering:
            model = pipeline("question-answering", model=model_name, tokenizer=model_name, device=device)
        else:
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForCausalLM.from_pretrained(model_name)
            model.to(device)

        # Append the new user message to the chat history
        chat_history.append({"role": "user", "content": message})

        if is_question_answering:
            if rag:
                filtered_contexts = extract_and_filter_contexts(data, min_length=100, max_length=4000)
                context = generate_rag_context(message, filtered_contexts, selected_context_window)
            else:
                context = "\n".join([turn["content"] for turn in chat_history if turn["role"] == "user"])

            if debug:
                print("context:\n" + context)
                print("message:\n" + message)

            # Call the pipeline for question-answering
            answer = model(question=message, context=context)
            response = answer['answer']

        else:
            # Prepare the conversation history for a regular chatbot
            conversation = ""
            for turn in chat_history:
                if turn["role"] == "user":
                    conversation += f"User: {turn['content']}\n"
                else:
                    conversation += f"Assistant: {turn['content']}\n"

            if debug:
                print("Conversation being sent to the model:\n", conversation)

            # Encode the input and generate a response
            inputs = tokenizer.encode(conversation + "Assistant:", return_tensors='pt').to(device)
            outputs = model.generate(
                inputs,
                max_length=inputs.shape[1] + 100,
                pad_token_id=tokenizer.eos_token_id,
                do_sample=True,
                top_p=0.95,
                top_k=50,
                temperature=0.7,
                eos_token_id=tokenizer.eos_token_id,
            )
            response = tokenizer.decode(outputs[0], skip_special_tokens=True)

            # Extract the assistant's reply
            response = response[len(conversation):].strip()
            if "User:" in response:
                response = response.split("User:")[0].strip()

        # Append the assistant's response to the chat history
        chat_history.append({"role": "assistant", "content": response})
        if debug:
            print("Generated response:", response)
            print("Configurations:")
            print(f"Model Name: {model_name}")
            print(f"Is Question Answering: {is_question_answering}")
            print(f"RAG Enabled: {rag}")
            print(f"Selected Context Window: {selected_context_window}")

        # Return the updated chat history and the assistant's response
        display_history = [[turn["content"], chat_history[i + 1]["content"]] for i, turn in enumerate(chat_history[:-1]) if turn["role"] == "user" and i + 1 < len(chat_history)]
        return display_history, chat_history

    except Exception as e:
        # Capture the traceback details
        error_message = f"An error occurred: {str(e)}"
        detailed_error = traceback.format_exc()
        chat_history.append({"role": "assistant", "content": error_message})
        if debug:
            print("Error Details:\n", detailed_error)

        # Ensure safe generation of the display history
        try:
            display_history = [[turn["content"], chat_history[i + 1]["content"]] for i, turn in enumerate(chat_history[:-1]) if turn["role"] == "user" and i + 1 < len(chat_history)]
        except Exception as history_error:
            if debug:
                print("Error while generating display history:", str(history_error))
            display_history = []

        return display_history, chat_history

# Gradio Interface Configuration
def run_prod_chatbot(local=True):
    with gr.Blocks() as demo:
        gr.Markdown("""
        <div style="text-align: center;">
            <h1><strong>SQuAD Q&A ChatBot</strong></h1>
            <h3>Authors: <a href="https://github.com/zainnobody">Zain Ali</a> & <a href="https://github.com/AIBenHopwood/">Ben Hopwood</a></h3>
            <p>
                <a href="https://github.com/zainnobody/AAI-520-Final-Project" target="_blank">Code: GitHub link</a> &nbsp;|&nbsp;
                <a href="https://huggingface.co/zainnobody/AAI-520-Final-Project-Models" target="_blank">Models: Huggingface link</a>
            </p>
        </div>
        
        <div style="text-align: center;">
            <p>
                This project aims to develop a chatbot capable of multi-turn, context-adaptive conversations across various topics, using the Stanford Question Answering Dataset (SQuAD) as the primary source for training.
            </p>
        </div>
        
        <div style="text-align: center;">
            <h4>University of San Diego - AAI 520</h4>
        </div>
        
                """)
        with gr.Row(variant="compact"):
            model_dropdown = gr.Dropdown(
                choices=electra_models + other_models,
                label="Select Model",
                value="./models/fine_tuned_electra_model_all"
            )
            # Column for Use RAG and Debug Mode checkboxes
            with gr.Column():
                rag_checkbox = gr.Checkbox(
                    label="Use RAG", 
                    value=True, 
                    interactive=True
                )
                debug_checkbox = gr.Checkbox(
                    label="Debug Mode", 
                    value=False
                )
            context_window_dropdown = gr.Dropdown(
                choices=[1, 2, 3],
                label="Select Context Window",
                value=1
            )
    
        # Commented out the is_question_answering_checkbox, making it auto detectable. Leaving this as a reminder that other models do not use pipeline
        # is_question_answering_checkbox = gr.Checkbox(
        #     label="Use Question Answering (Electra Only)", 
        #     value=True
        # )
        
        chatbot = gr.Chatbot()
        state = gr.State([])
        
        with gr.Row():
            # Textbox taking 75% of the space
            msg = gr.Textbox(label="Your message", placeholder="Type your message here and press Enter", scale=3)
            # Send button taking 25% of the space and stretching full width
            send_btn = gr.Button("Send", scale=1)

        
            
        send_btn.click(lambda message, chat_history, model_name, debug, rag, selected_context_window: generate_response(message, chat_history, model_name, debug, rag, selected_context_window),
                      inputs=[msg, state, model_dropdown, debug_checkbox, rag_checkbox, context_window_dropdown],
                      outputs=[chatbot, state])
        msg.submit(lambda message, chat_history, model_name, debug, rag, selected_context_window: generate_response(message, chat_history, model_name, debug, rag, selected_context_window),
                   inputs=[msg, state, model_dropdown, debug_checkbox, rag_checkbox, context_window_dropdown],
                   outputs=[chatbot, state])

    if local:
        demo.launch(share=True)
    else:
        demo.launch(server_name="0.0.0.0", server_port=None)

# Launch the Gradio app
run_prod_chatbot()