import pandas as pd
import gradio as gr
import pyterrier as pt
pt.init()
from pyterrier_gradio import Demo, MarkdownFile, interface, df2code, code2md, EX_R
from pyterrier_t5 import MonoT5ReRanker

model = MonoT5ReRanker()

COLAB_NAME = 'pyterrier_t5.ipynb'
COLAB_INSTALL = '''
!pip install -q git+https://github.com/terrier-org/pyterrier_t5
'''.strip()

def predict(input):
  code = f'''import pandas as pd
import pyterrier as pt ; pt.init()
from pyterrier_t5 import MonoT5ReRanker

model = MonoT5ReRanker()

model({df2code(input)})
'''
  res = model(input)
  res['score'] = res['score'].map(lambda x: round(x, 4))
  res = res.sort_values(['qid', 'rank'])
  return (res, code2md(code, COLAB_INSTALL, COLAB_NAME, colab=False))

interface(
  MarkdownFile('README.md'),
  Demo(
    predict,
    EX_R,
    []
  ),
  MarkdownFile('wrapup.md'),
).launch(share=False)