app.py

import streamlit as st
from datasets import load_dataset
from langchain.chains import LLMChain
from langchain.llms import HuggingFaceHub
from langchain.prompts import PromptTemplate
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 HuggingFaceHub model (T5)
llm = HuggingFaceHub(repo_id="google/pegasus-xsum", model_kwargs={"temperature": 0.7})

# 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 (via Langchain)")
input_text = st.text_area("📝 Paste your conversation:")

if st.button("Generate Summary"):
    if input_text.strip():
        # Create prompt using FewShotChatMessagePromptTemplate
        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```")

        # Create the prompt chain
        prompt_template = PromptTemplate(input_variables=["input"], template="{input}")
        chain = LLMChain(llm=llm, prompt=prompt_template)

        # Get the summary from the model
        summary = chain.run(input_text)
        st.success("✅ Summary:")
        st.write(summary)
    else:
        st.warning("Please enter some text.")