app.py

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import pickle
import torch


import io

class CPU_Unpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module == 'torch.storage' and name == '_load_from_bytes':
            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
        else:
            return super().find_class(module, name)

#contents = pickle.load(f) becomes...
#contents = CPU_Unpickler(f).load()


checkpoint = "my_t5.sav"

#load model from drive
with open(checkpoint, "rb") as f:
    model=  CPU_Unpickler(f).load()
    
#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"))