import torch

import gradio as gr

from transformers import T5ForConditionalGeneration, T5Tokenizer

import os

#import whisper

import matplotlib as plt

# whisper_model = whisper.load_model('large-v2')      # Whisper 모델을 불러오기


path = "Hyeonsieun/NTtoGT_7epoch"
tokenizer = T5Tokenizer.from_pretrained(path)
model = T5ForConditionalGeneration.from_pretrained(path)

def do_correction(text):
    input_text = f"translate the text pronouncing the formula to a LaTeX equation: {text}"
    inputs = tokenizer.encode(
        input_text,
        return_tensors='pt',
        max_length=325,
        padding='max_length',
        truncation=True
    )

    # Get correct sentence ids.
    corrected_ids = model.generate(
        inputs,
        max_length=325,
        num_beams=5, # `num_beams=1` indicated temperature sampling.
        early_stopping=True
    )

    # Decode.
    corrected_sentence = tokenizer.decode(
        corrected_ids[0],
        skip_special_tokens=False
    )
    return corrected_sentence






gr.Interface(fn=do_correction, inputs="text", outputs="text").launch()