app.py

import streamlit as st
from datasets import load_dataset
from transformers import pipeline
from langchain.prompts.few_shot import FewShotChatMessagePromptTemplate
from langchain.prompts.example_selector import LengthBasedExampleSelector

# Load dataset (small subset)
@st.cache_data
def load_examples():
    dataset = load_dataset("knkarthick/dialogsum", split="train[:5]")  # Take only 5 for speed
    examples = []
    for example in dataset:
        examples.append({
            "input": example["dialogue"],
            "output": example["summary"]
        })
    return examples

examples = load_examples()

# Set up the summarization model
summarizer = pipeline("summarization", model="t5-small")

# Few-shot prompt template
example_prompt = FewShotChatMessagePromptTemplate.from_examples(
    examples=examples,
    example_selector=LengthBasedExampleSelector(examples=examples, max_length=1000),
    input_variables=["input"],
    prefix="You are a helpful assistant that summarizes dialogues. Examples:",
    suffix="Now summarize this:\n{input}"
)

# Streamlit UI
st.title("💬 Dialogue Summarizer using Few-Shot Prompt + T5")
input_text = st.text_area("📝 Paste your conversation:")

if st.button("Generate Summary"):
    if input_text.strip():
        # Create prompt
        messages = example_prompt.format_messages(input=input_text)

        with st.expander("📋 Generated Prompt"):
            for msg in messages:
                st.markdown(f"**{msg.type.upper()}**:\n```\n{msg.content}\n```")

        # Generate summary using model
        full_prompt = "\n".join([m.content for m in messages if m.type == "human"])
        summary = summarizer("summarize: " + input_text, max_length=80, min_length=15, do_sample=False)[0]['summary_text']
        st.success("✅ Summary:")
        st.write(summary)
    else:
        st.warning("Please enter some text.")