import gradio as gr
from optimum.onnxruntime import ORTModelForSeq2SeqLM
from transformers import pipeline
import logging
import os

logger = logging.Logger('sc_summ')
logger.setLevel(logging.INFO)
# File handler (for logging to file)
local_dir = os.path.dirname(__file__)
file_handler = logging.FileHandler(os.path.join(local_dir, "sc_summ.log"), encoding='utf-8')
file_handler.setLevel(logging.INFO)
file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
file_handler.setFormatter(file_formatter)
logger.addHandler(file_handler)  

def get_summarizer():
    model = ORTModelForSeq2SeqLM.from_pretrained(
        "shorecode/t5-efficient-tiny-summarizer-general-purpose-v3"
    )        
    return pipeline(
        "summarization",
        model=model,
        tokenizer="shorecode/t5-efficient-tiny-summarizer-general-purpose-v3",
    )

SUMMARY_MODEL = get_summarizer()

def summarize(source):
    try:
        summary = SUMMARY_MODEL(
            source,
            max_new_tokens=4000,
            min_length=20,
            no_repeat_ngram_size=2,
            num_beams=3, 
        )
        try:
            summary = summary[0]['summary_text']
        except (IndexError, KeyError):
            pass
        return summary
        return summary
    except Exception as e:
        logger.error(f"An error occured while summarizing: {e}")
        return "An error occurred"

def greet(name):
    return "Hello " + __file__ + "!!"

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