from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

checkpoint = "Shivam29rathore/t5_10k_base"

#tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)


def summarize(word):
    import os
    data_path = "/tmp/"
    if not os.path.exists(data_path):
        os.makedirs(data_path)
    input_ = "/tmp/input.txt"

    with open(input_, "w") as file:
        file.write(word)
    # read the written txt into a variable
    with open(input_ , 'r') as f:
        text_ = f.read()

    def clean_data(texts):
        import re
        words = list()
        for text in texts.split():
            text = re.sub(r'\n','',text)
            text = re.sub(r'\s$','',text)
            words.append(text)

        return "summarize " + " ".join(words)
    text = clean_data(text_)

    final_summary = []
    for x in range(0,len(text)-1256,1256):
        text_to_summarize= text[x:x+1256]
        final_summary.append(model.predict(text_to_summarize))

    final_list = list(itertools.chain.from_iterable(final_summary))
    final_list = ''.join(final_list)
    return final_list
    
    
import gradio as gr

iface = gr.Interface(fn= summarize, 
                     inputs =gr.inputs.Textbox(lines=15,placeholder="Enter your text !!"),
                     outputs="text",title="Document Summarizer",description ="An AI that makes your life easier by helping you summarise long texts.")
iface.launch(auth=("docai","ailabs"))