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

# Load few-shot examples from dialogsum
@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()

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

# Few-shot prompt template (no prefix/suffix here)
few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples
)

# Now add intro system message + user input separately
final_prompt = ChatPromptTemplate.from_messages([
    SystemMessage(content="The following are examples of dialogues and their summaries."),
    *few_shot_prompt.get_messages(),  # Use get_messages() instead of .messages
    ("human", "Summarize the following dialog:\n\n{dialogue}")
])

# Load Pegasus model from HF inference API
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:
    # Format messages
    messages = final_prompt.format_messages(dialogue=user_input)
    
    # Get response
    response = llm(messages)
    
    # Output
    st.subheader("📌 Summary:")
    st.write(response)