import gradio as gr

from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    AutoConfig,
    pipeline,
)

model_name = "sagard21/python-code-explainer"

tokenizer = AutoTokenizer.from_pretrained(model_name, padding=True)

model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

config = AutoConfig.from_pretrained(model_name)

model.eval()

pipe = pipeline("summarization", model=model_name, config=config, tokenizer=tokenizer)

def generate_text(text_prompt):
  response = pipe(text_prompt)
  return response[0]['summary_text']

textbox1 = gr.Textbox(value = """
class Solution(object):
    def isValid(self, s):
        stack = []
        mapping = {")": "(", "}": "{", "]": "["}
        for char in s:
            if char in mapping:
                top_element = stack.pop() if stack else '#'
                if mapping[char] != top_element:
                    return False
            else:
                stack.append(char)
        return not stack""")

textbox2 = gr.Textbox()

image_code = gr.Image(value = "https://huggingface.co/spaces/PinoCorgi/CodeExplainerPython/blob/main/code.jpg")
image_output = gr.Image(value = "https://huggingface.co/spaces/PinoCorgi/CodeExplainerPython/blob/main/output.jpg")

demo = gr.Interface(fn = generate_text, inputs = textbox1, outputs = textbox2)

image_code
image_output

if __name__ == "__main__":
    print(f"The Inference will take approximately 1 min 30 Seconds. Attached are example images of input and output.")
    demo.launch()