import gradio as gr
from transformers import pipeline
model_id = "knkarthick/MEETING-SUMMARY-BART-LARGE-XSUM-SAMSUM-DIALOGSUM-AMI"

generator = pipeline(task="text2text-generation", model=model_id)

def split_paragraph(paragraph, max_chunk_size=1024):
    words = paragraph.split()
    chunks = []
    current_chunk = []
    current_chunk_size = 0

    for word in words:
        word_len = len(word) + 1  # Add 1 for the space
        if current_chunk_size + word_len <= max_chunk_size:
            current_chunk.append(word)
            current_chunk_size += word_len
        else:
            chunks.append(' '.join(current_chunk))
            current_chunk = [word]
            current_chunk_size = word_len

    if current_chunk:
        chunks.append(' '.join(current_chunk))

    return chunks

def launch(input):
    if len(input) > 1024:
        return " ".join([res["generated_text"] for res in generator(split_paragraph(input))])
    return generator(input)[0]["generated_text"]

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