app.py

import streamlit as st
from datasets import load_dataset
from langchain.llms import HuggingFaceEndpoint
from langchain.prompts import FewShotChatMessagePromptTemplate, ChatPromptTemplate

# Load dataset from HuggingFace
@st.cache_data
def load_examples(n=3):
    dataset = load_dataset("knkarthick/dialogsum", split="train[:20]")
    return [{"dialogue": row["dialogue"], "summary": row["summary"]} for row in dataset.select(range(n))]

examples = load_examples()

# Format examples
example_prompt = ChatPromptTemplate.from_messages([
    ("human", "Summarize the following dialog:\n\n{dialogue}"),
    ("ai", "{summary}")
])

# Few-shot setup
few_shot_prompt = FewShotChatMessagePromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    suffix="Summarize the following dialog:\n\n{dialogue}",
    input_variables=["dialogue"],
    prefix="The following are examples of dialogues and their summaries."
)

# Load HF summarizer model (Pegasus)
llm = HuggingFaceEndpoint(
    repo_id="google/pegasus-xsum",
    task="text2text-generation",
    model_kwargs={"temperature": 0.3, "max_new_tokens": 128}
)

# Streamlit UI
st.set_page_config(page_title="DialogSum Few-Shot Summarizer", page_icon="🧠")
st.title("🧠 Few-Shot Dialog Summarizer")
st.markdown("Uses real examples from `dialogsum` to guide the summary output.")

user_input = st.text_area("✍️ Paste your dialogue here:", height=200)

if user_input:
    messages = few_shot_prompt.format_messages(dialogue=user_input)
    response = llm(messages)
    st.subheader("📌 Summary:")
    st.write(response)