import os
import io
from IPython.display import Image, display, HTML
from PIL import Image
import base64
from transformers import pipeline, AutoTokenizer

# Initialize summarizer and tokenizer
summarizer = pipeline ("summarization", model="sshleifer/distilbart-cnn-12-6", tokenizer="sshleifer/distilbart-cnn-12-6")
tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-cnn-12-6")

import json

def summarize_text(input_text):
    """Summarizes the given input text.

    Args:
        input_text (str): The text to be summarized.

    Returns:
        dict: A dictionary containing the summary under the 'summary' key.
    """

    # Tokenize and truncate input if necessary
    max_length = tokenizer.model_max_length
    inputs = tokenizer(input_text, truncation=True, max_length=max_length, return_tensors="pt")

    # Generate summary
    summary_ids = summarizer.model.generate(inputs.input_ids, max_length=50, min_length=10, do_sample=False)
    summary_text = tokenizer.decode(summary_ids[0], skip_special_tokens=True)

    # Return summary as a dictionary
    return {"summary": summary_text}

def generate_summary(input):
    output = summarize_text(input)
    return output
    
gr.close_all()
demo = gr.Interface(fn=generate_summary, 
                        inputs=[gr.Textbox(label="Text to summarize", lines=6)], 
                        outputs=[gr.Textbox(label="Summary", lines=3)],
                        title="Text Summarization",
                        description="Summarize text using the 'shleifer/distilbart-cnn-12-6' language model.",                    
                    )
demo.launch(share=True, server_port=int(os.environ['PORT']))