from transformers import EncoderDecoderModel,GPT2LMHeadModel,BertTokenizer
from datasets import load_dataset
import gradio as gr
import pandas as pd
import torch


tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
Model = EncoderDecoderModel.from_pretrained('damilojohn/Bert2BertForTextDescrambling')
def descramble(prompt):
  input = tokenizer(prompt,return_tensors='pt')
  input_id = input.input_ids
  attention_mask = input.attention_mask
  max_length = len(prompt.split(' '))
  output = Model.generate(input_ids=input_id,attention_mask=attention_mask,)
  output = tokenizer.decode(output[0],skip_special_tokens=True)
  return gr.Textbox.update(value=output)
 
def set_example(example):
    return gr.TextArea.update(value=example[0])
  
 demo = gr.Blocks()
with demo:
    gr.Markdown(
        '''
        # A Text Descrambler 😎😎
        Turn your Incoherent Sentences to Grammatically correct Sentences. 
        This was built using transformers and Gradio 
        ''')
    with gr.Row():
        with gr.Column():
            gr.Markdown(
                ''' 
                Enter a meaningless sentence here
            
                ''')
            prompt = gr.TextArea(
                value = examples[0][0],
                placeholder = "Enter A Text to see it's correct form "
            )
            example_prompts = gr.Dataset(
                components = [prompt],
                samples = examples)
        with gr.Column():
            find_answer = gr.Button('Click here to generate your sentence 👀🤺').style(full_width=False)
        with gr.Column():
            answer = gr.Textbox(label='Answer',placeholder = "Correct Form")
        with gr.Column():
          gr.Markdown(
              '''
              ## Still Under Construction 💀🎃⏳, 
              anything you see take it like that 
              ''')
    find_answer.click(
        fn=descramble,
        inputs=[prompt],
        outputs=[answer]
    )
    example_prompts.click(
        fn=set_example,
        inputs=[example_prompts],
        outputs=example_prompts.components,
    )

    
        
demo.launch()