from datasets import load_dataset
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr

def compute_similarity(dataset_name):
    # Load dataset
    dataset = load_dataset(dataset_name)
    # Dummy similarity computation (replace with your metric)
    data = np.random.rand(10, 10)
    # Create heatmap
    fig, ax = plt.subplots()
    cax = ax.matshow(data, cmap='viridis')
    plt.colorbar(cax)
    return fig

with gr.Blocks() as demo:
    dataset_name = gr.Textbox(label="Enter Dataset Name (e.g., 'imdb')")
    heatmap_plot = gr.Plot(label="Similarity Heatmap")
    compute_button = gr.Button("Compute Similarity")

    compute_button.click(
        fn=compute_similarity,
        inputs=dataset_name,
        outputs=heatmap_plot
    )

demo.launch()