import gradio as gr
import pyterrier as pt
pt.init()
from pyterrier_doc2query import Doc2Query

doc2query = Doc2Query()

def predict(input_df):
  return input_df

gr.Interface(
    predict,
    inputs=gr.Dataframe(
            headers=["docno", "text"],
            datatype=["str", "str"],
            row_count=5,
            col_count=(2, "fixed"),
        ),
    outputs="dataframe",
    title="pyterrier_doc2query",
).launch()