import gradio as gr
import pandas as pd

def topn_tokens(sequence, domain_bounds, n):
    example_dict = {}
    chars = list(sequence)
    start_index = domain_bounds['start'][0] - 1  
    end_index = domain_bounds['end'][0] - 1      

    for i in range(len(sequence)):
        if start_index <= i <= end_index:
            example_dict[chars[i]] = 'yo'
    
    df = pd.DataFrame(list(example_dict.items()), columns=['Original Residue', 'Predicted Residues'])
    return df

demo = gr.Interface(
    fn=topn_tokens,
    inputs=[
        "text",
        gr.Dataframe(
            headers=["start", "end"],
            datatype=["number", "number"],
            row_count=(1, "fixed"),
            col_count=(2, "fixed"),
        ),
        gr.Dropdown([str(i) for i in range(1, 21)]),  # Dropdown with numbers from 1 to 20
    ],
    outputs="dataframe",
    description="Choose a number between 1-20 to predict n tokens for each position. Choose the start and end index of the domain of interest (indexing starts at 1).",
)

if __name__ == "__main__":
    demo.launch()