UI/gradio.py

import gradio as gr
from utils.watermark import Watermarker
from utils.config import load_config
from renderers.highlighter import highlight_common_words, highlight_common_words_dict, reparaphrased_sentences_html
from renderers.tree import generate_subplot1, generate_subplot2
from pathlib import Path
import time
from typing import Dict, List, Tuple, Any
import plotly.graph_objects as go

class WatermarkerInterface:
    def __init__(self, config):
        self.pipeline = Watermarker(config)    
        self.common_grams = {}
        self.highlight_info = []    
        self.masked_sentences = []
        
        # Add tracking dictionaries for indexing
        self.masked_sentence_indices = {}  # Maps original sentences to masked indices
        self.sampled_sentence_indices = {}  # Maps masked sentences to sampling indices
        self.reparaphrased_indices = {}    # Maps sampled sentences to reparaphrased indices
        
    def handle_paraphrase(self, prompt: str) -> Tuple[str, str, str, str]:
        """Wrapper for paraphrasing that includes highlighting"""
        start_time = time.time()
        
        # Run paraphrasing
        self.pipeline.Paraphrase(prompt)

        # Step 1: Process the original sentence first
        seen_ngrams = {}  # Stores first occurrence index of each n-gram
        original_indexed_ngrams = []  # Final indexed list for original
        
        original_sentence = self.pipeline.user_prompt
        original_ngrams = self.pipeline.common_grams.get(original_sentence, {})

        # Step 1.1: Extract n-grams and their first occurrence index
        ngram_occurrences = [
            (min(indices, key=lambda x: x[0])[0], gram)  # Get first index
            for gram, indices in original_ngrams.items()
        ]

        # Step 1.2: Sort n-grams based on their first occurrence
        ngram_occurrences.sort()

        # Step 1.3: Assign sequential indices
        for idx, (position, gram) in enumerate(ngram_occurrences, start=1):
            seen_ngrams[gram] = idx  # Assign sequential index
            original_indexed_ngrams.append((idx, gram))

        print("Original Indexed N-grams:", original_indexed_ngrams)

        #generate highlight_info
        colors = ["red", "blue", "purple", "green", "orange"]
        highlight_info = [
            (ngram, colors[i % len(colors)])
            for i, (index, ngram) in enumerate(original_indexed_ngrams)
        ]
        common_grams = original_indexed_ngrams
        self.highlight_info = highlight_info
        self.common_grams = common_grams

        # Step 2: Process paraphrased sentences and match indices
        paraphrase_indexed_ngrams = {}

        for sentence in self.pipeline.paraphrased_sentences:
            sentence_ngrams = []  # Stores n-grams for this sentence
            sentence_ngrams_dict = self.pipeline.common_grams.get(sentence, {})

            for gram, indices in sentence_ngrams_dict.items():
                first_occurrence = min(indices, key=lambda x: x[0])[0]

                # Use the original's index if exists, otherwise assign a new one
                if gram in seen_ngrams:
                    index = seen_ngrams[gram]  # Use the same index as original
                else:
                    index = len(seen_ngrams) + 1  # Assign new index
                    seen_ngrams[gram] = index  # Store it

                sentence_ngrams.append((index, gram))

            sentence_ngrams.sort()
            paraphrase_indexed_ngrams[sentence] = sentence_ngrams

        print("Paraphrase Indexed N-grams:", paraphrase_indexed_ngrams)

        # Step 3: Generate highlighted versions using the renderer
        highlighted_prompt = highlight_common_words(
            common_grams, 
            [self.pipeline.user_prompt], 
            "Original Prompt with Highlighted Common Sequences"
        )

        highlighted_accepted = highlight_common_words_dict(
            common_grams,
            self.pipeline.selected_sentences,
            "Accepted Paraphrased Sentences with Entailment Scores"
        )

        highlighted_discarded = highlight_common_words_dict(
            common_grams,
            self.pipeline.discarded_sentences,
            "Discarded Paraphrased Sentences with Entailment Scores"
        )

        execution_time = f"<div class='execution-time'>Step 1 completed in {time.time() - start_time:.2f} seconds</div>"

        return highlighted_prompt, highlighted_accepted, highlighted_discarded, execution_time

    def handle_masking(self):
        start_time = time.time()
        masking_results = self.pipeline.Masking()
        trees = []
        highlight_info = self.highlight_info
        common_grams = self.common_grams
        sentence_to_masked = {}
        self.masked_sentence_indices = {}

        for strategy, sentence_dict in masking_results.items():
            for sent, data in sentence_dict.items():
                if sent not in sentence_to_masked:
                    sentence_to_masked[sent] = []
                masked_sentence = data.get("masked_sentence", "")
                if masked_sentence:
                    sentence_to_masked[sent].append((masked_sentence, strategy))

        plot_idx = 1
        for original_sentence, masked_sentences_data in sentence_to_masked.items():
            if not masked_sentences_data:
                continue
            masked_idx = 1
            for masked_sentence, strategy in masked_sentences_data:
                index = f"{plot_idx}{masked_idx}"
                if original_sentence not in self.masked_sentence_indices:
                    self.masked_sentence_indices[original_sentence] = {}
                key = f"{strategy}_{masked_sentence}"
                self.masked_sentence_indices[original_sentence][key] = {
                    'index': index,
                    'strategy': strategy,
                    'masked_sentence': masked_sentence
                }
                masked_idx += 1

            masked_sentences = [ms[0] for ms in masked_sentences_data]
            indexed_masked_sentences = []
            verified_strategies = []
            for masked_sentence, strategy in masked_sentences_data:
                key = f"{strategy}_{masked_sentence}"
                entry = self.masked_sentence_indices[original_sentence][key]
                idx = entry['index']
                indexed_masked_sentences.append(f"[{idx}] {masked_sentence}")
                verified_strategies.append(entry['strategy'])

            try:
                fig = generate_subplot1(
                    original_sentence, 
                    indexed_masked_sentences,
                    verified_strategies,
                    highlight_info,
                    common_grams
                )
                trees.append(fig)
            except Exception as e:
                print(f"Error generating plot: {e}")
                trees.append(go.Figure())
            plot_idx += 1

        while len(trees) < 10:
            trees.append(go.Figure())

        execution_time = f"<div class='execution-time'>Step 2 completed in {time.time() - start_time:.2f} seconds</div>"
        return trees[:10] + [execution_time]

    def handle_sampling(self) -> Tuple[List[go.Figure], str]:
        start_time = time.time()
        sampling_results = self.pipeline.Sampling()
        trees = []
        self.sampled_sentence_indices = {}
        organized_results = {}

        for sampling_strategy, masking_dict in sampling_results.items():
            for masking_strategy, sentences in masking_dict.items():
                for original_sentence, data in sentences.items():
                    if original_sentence not in organized_results:
                        organized_results[original_sentence] = {}
                    if masking_strategy not in organized_results[original_sentence]:
                        organized_results[original_sentence][masking_strategy] = {
                            "masked_sentence": data.get("masked_sentence", ""),
                            "sampled_sentences": {}
                        }
                    organized_results[original_sentence][masking_strategy]["sampled_sentences"][sampling_strategy] = data.get("sampled_sentence", "")

        plot_idx = 1
        for original_sentence, data in organized_results.items():
            masked_sentences = []
            all_sampled_sentences = []
            indexed_sampled_sentences = []
            masked_indices = self.masked_sentence_indices.get(original_sentence, {})

            for masking_strategy, masking_data in list(data.items())[:3]:
                masked_sentence = masking_data.get("masked_sentence", "")
                if masked_sentence:
                    masked_sentences.append(masked_sentence)
                    masked_idx = None
                    for ms_key, ms_data in masked_indices.items():
                        if ms_key == f"{masking_strategy}_{masked_sentence}":
                            masked_idx = ms_data['index']
                            break

                    if not masked_idx:
                        print(f"Warning: No index found for masked sentence: {masked_sentence}")
                        continue

                    sample_count = 1
                    for sampling_strategy, sampled_sentence in masking_data.get("sampled_sentences", {}).items():
                        if sampled_sentence:
                            sample_idx = f"{masked_idx}.{sample_count}"
                            if masked_sentence not in self.sampled_sentence_indices:
                                self.sampled_sentence_indices[masked_sentence] = {}
                            self.sampled_sentence_indices[masked_sentence][sampled_sentence] = {
                                'index': sample_idx,
                                'strategy': sampling_strategy
                            }
                            indexed_sampled_sentences.append(f"[{sample_idx}] {sampled_sentence}")
                            all_sampled_sentences.append(sampled_sentence)
                            sample_count += 1

            if masked_sentences:
                indexed_masked_sentences = []
                for ms in masked_sentences:
                    idx = ""
                    for ms_key, ms_data in masked_indices.items():
                        if ms_key.endswith(f"_{ms}"):
                            idx = ms_data['index']
                            break
                    indexed_masked_sentences.append(f"[{idx}] {ms}")

                try:
                    fig = generate_subplot2(
                        indexed_masked_sentences,
                        indexed_sampled_sentences,
                        self.highlight_info,
                        self.common_grams
                    )
                    trees.append(fig)
                except Exception as e:
                    print(f"Error generating subplot for {original_sentence}: {e}")
                    trees.append(go.Figure())
            plot_idx += 1

        print("Sampled sentence indices:", self.sampled_sentence_indices)

        while len(trees) < 10:
            trees.append(go.Figure())

        execution_time = f"<div class='execution-time'>Step 3 completed in {time.time() - start_time:.2f} seconds</div>"

        return trees[:10] + [execution_time]

    def handle_reparaphrasing(self) -> Tuple[List[str], str]:
        start_time = time.time()
        results = self.pipeline.re_paraphrasing()
        html_outputs = []
        self.reparaphrased_indices = {}
        tab_count = 1

        for sampling_strategy, masking_dict in results.items():
            for masking_strategy, sentences in masking_dict.items():
                for original_sent, data in sentences.items():
                    sampled_sentence = data.get("sampled_sentence", "")
                    if not sampled_sentence or not data["re_paraphrased_sentences"]:
                        continue

                    sampled_index = None
                    for masked_sent, sampled_dict in self.sampled_sentence_indices.items():
                        if sampled_sentence in sampled_dict:
                            sampled_index = sampled_dict[sampled_sentence]['index']
                            break

                    if not sampled_index:
                        sampled_index = "unknown"

                    indexed_reparaphrased = []
                    for i, rp_sent in enumerate(data["re_paraphrased_sentences"], 1):
                        rp_idx = f"{tab_count}.({sampled_index}).{i}"
                        if sampled_sentence not in self.reparaphrased_indices:
                            self.reparaphrased_indices[sampled_sentence] = {}
                        self.reparaphrased_indices[sampled_sentence][rp_sent] = rp_idx
                        indexed_reparaphrased.append(f"[{rp_idx}] {rp_sent}")

                    print(f"Reparaphrasing {tab_count}.({sampled_index}): {' '.join(sampled_sentence.split()[:5])}...")
                    html = reparaphrased_sentences_html(indexed_reparaphrased)
                    html_outputs.append(html)
                    tab_count += 1

        print("Reparaphrased indices:", self.reparaphrased_indices)

        while len(html_outputs) < 150:
            html_outputs.append("")

        execution_time = f"<div class='execution-time'>Step 4 completed in {time.time() - start_time:.2f} seconds</div>"

        return html_outputs[:150] + [execution_time]
    
def create_gradio_interface(config):
    """Creates the Gradio interface with the updated pipeline"""
    interface = WatermarkerInterface(config)
    
    with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
        #CSS to enable scrolling for reparaphrased sentences and sampling plots
        demo.css = """
/* Set fixed height for the reparaphrased tabs container only */
.gradio-container .tabs[id="reparaphrased-tabs"],
.gradio-container .tabs[id="sampling-tabs"] {
    overflow-x: hidden;
    white-space: normal;
    border-radius: 8px;
    max-height: 600px; /* Set fixed height for the entire tabs component */
    overflow-y: auto; /* Enable vertical scrolling inside the container */
}

/* Tab content styling for reparaphrased and sampling tabs */
.gradio-container .tabs[id="reparaphrased-tabs"] .tabitem,
.gradio-container .tabs[id="sampling-tabs"] .tabitem {
    overflow-x: hidden;
    white-space: normal;
    display: block;
    border-radius: 8px;
}

/* Make the tab navigation fixed at the top for scrollable tabs */
.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav {
    display: flex;
    overflow-x: auto;
    white-space: nowrap;
    scrollbar-width: thin;
    border-radius: 8px;
    scrollbar-color: #888 #f1f1f1;
    position: sticky;
    top: 0;
    background: white;
    z-index: 100;
}

/* Dropdown menu for scrollable tabs styling */
.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav .tab-dropdown,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav .tab-dropdown {
    position: relative;
    display: inline-block;
}

.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav .tab-dropdown-content,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav .tab-dropdown-content {
    display: none;
    position: absolute;
    background-color: #f9f9f9;
    min-width: 160px;
    box-shadow: 0px 8px 16px 0px rgba(0,0,0,0.2);
    z-index: 1;
    max-height: 300px;
    overflow-y: auto;
}

.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav .tab-dropdown:hover .tab-dropdown-content,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav .tab-dropdown:hover .tab-dropdown-content {
    display: block;
}

/* Scrollbar styling for scrollable tabs */
.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav::-webkit-scrollbar,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav::-webkit-scrollbar {
    height: 8px;
    border-radius: 8px;
}

.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav::-webkit-scrollbar-track,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav::-webkit-scrollbar-track {
    background: #f1f1f1;
    border-radius: 8px;
}

.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav::-webkit-scrollbar-thumb,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav::-webkit-scrollbar-thumb {
    background: #888;
    border-radius: 8px;
}

/* Tab button styling for scrollable tabs */
.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav .tab-item,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav .tab-item {
    flex: 0 0 auto;
    border-radius: 8px;
}

/* Plot container styling specifically for sampling tabs */
.gradio-container .tabs[id="sampling-tabs"] .plot-container {
    min-height: 600px;
    max-height: 1800px;
    overflow-y: auto;
}

/* Ensure text wraps in HTML components */
.gradio-container .prose {
    white-space: normal;
    word-wrap: break-word;
    overflow-wrap: break-word;
}

/* Dropdown button styling for scrollable tabs */
.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav .tab-dropdown button,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav .tab-dropdown button {
    background-color: #f0f0f0;
    border: 1px solid #ddd;
    border-radius: 4px;
    padding: 5px 10px;
    cursor: pointer;
    margin: 2px;
}

.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav .tab-dropdown button:hover,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav .tab-dropdown button:hover {
    background-color: #e0e0e0;
}

/* Style dropdown content items for scrollable tabs */
.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav .tab-dropdown-content div,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav .tab-dropdown-content div {
    padding: 8px 12px;
    cursor: pointer;
}

.gradio-container .tabs[id="reparaphrased-tabs"] .tab-nav .tab-dropdown-content div:hover,
.gradio-container .tabs[id="sampling-tabs"] .tab-nav .tab-dropdown-content div:hover {
    background-color: #e0e0e0;
}

/* Custom styling for execution time display */
.execution-time {
    text-align: right;
    padding: 8px 16px;
    font-family: inherit;
    color: #555;
    font-size: 0.9rem;
    font-style: italic;
    margin-left: auto;
    width: 100%;
    border-top: 1px solid #eee;
    margin-top: 8px;
}

/* Layout for section headers with execution time */
.section-header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    width: 100%;
    margin-bottom: 12px;
}

.section-header h3 {
    margin: 0;
}
"""
        gr.Markdown("# **AIISC Watermarking Model**")
        
        with gr.Column():
            gr.Markdown("## Input Prompt")
            user_input = gr.Textbox(
                label="Enter Your Prompt",
                placeholder="Type your text here..."
            )

        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown("## Step 1: Paraphrasing, LCS and Entailment Analysis")
            with gr.Column(scale=1):
                step1_time = gr.HTML()

        paraphrase_button = gr.Button("Generate Paraphrases")
        highlighted_user_prompt = gr.HTML(label="Highlighted User Prompt")

        with gr.Tabs():
            with gr.TabItem("Accepted Paraphrased Sentences"):
                highlighted_accepted_sentences = gr.HTML()
            with gr.TabItem("Discarded Paraphrased Sentences"):
                highlighted_discarded_sentences = gr.HTML()
            
        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown("## Step 2: Where to Mask?")
            with gr.Column(scale=1):
                step2_time = gr.HTML()

        masking_button = gr.Button("Apply Masking")
        gr.Markdown("### Masked Sentence Trees")
        tree1_plots = []
        with gr.Tabs() as tree1_tabs:
            for i in range(10):
                with gr.TabItem(f"Masked Sentence {i+1}"):
                    tree1 = gr.Plot()
                    tree1_plots.append(tree1)

        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown("## Step 3: How to Mask?")
            with gr.Column(scale=1):
                step3_time = gr.HTML()

        sampling_button = gr.Button("Sample Words")
        gr.Markdown("### Sampled Sentence Trees")

        tree2_plots = []
        # Add elem_id to make this tab container scrollable
        with gr.Tabs(elem_id="sampling-tabs") as tree2_tabs:
            for i in range(10):
                with gr.TabItem(f"Sampled Sentence {i+1}"):
                    # Add a custom class to the container to enable proper styling
                    with gr.Column(elem_classes=["plot-container"]):
                        tree2 = gr.Plot()
                        tree2_plots.append(tree2)
        
        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown("## Step 4: Re-paraphrasing")
            with gr.Column(scale=1):
                step4_time = gr.HTML()

        reparaphrase_button = gr.Button("Re-paraphrase")
        gr.Markdown("### Reparaphrased Sentences")
        reparaphrased_sentences_tabs = []
        with gr.Tabs(elem_id="reparaphrased-tabs") as reparaphrased_tabs:
            for i in range(150):
                with gr.TabItem(f"Reparaphrased Batch {i+1}"):
                    reparaphrased_sent_html = gr.HTML()
                    reparaphrased_sentences_tabs.append(reparaphrased_sent_html)
        
        # Connect the interface functions to the buttons
        paraphrase_button.click(
            interface.handle_paraphrase,
            inputs=user_input,
            outputs=[
                highlighted_user_prompt,
                highlighted_accepted_sentences,
                highlighted_discarded_sentences,
                step1_time
            ]
        )
        
        masking_button.click(
            interface.handle_masking,
            inputs=None,
            outputs=tree1_plots + [step2_time]
        )

        sampling_button.click(
            interface.handle_sampling,
            inputs=None,
            outputs=tree2_plots + [step3_time]
        )
        
        reparaphrase_button.click(
            interface.handle_reparaphrasing,
            inputs=None,
            outputs=reparaphrased_sentences_tabs + [step4_time]
        )
    
    return demo

if __name__ == "__main__":
    project_root = Path(__file__).parent.parent
    config_path = project_root / "utils" / "config.yaml"
    config = load_config(config_path)['PECCAVI_TEXT']
    
    create_gradio_interface(config).launch()