Update app.py

app.py

@@ -1,48 +1,47 @@
 import streamlit as st
-from langchain.prompts import FewShotChatMessagePromptTemplate
-from langchain.llms import HuggingFaceHub
 from datasets import load_dataset
+from langchain.llms import HuggingFaceEndpoint
+from langchain.prompts import FewShotChatMessagePromptTemplate, ChatPromptTemplate
 
-# Load the dataset
-dataset = load_dataset("knkarthick/dialogsum", split="train[:1%]")  # Load a small subset of the dataset for testing
+# 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))]
 
-# Extract the input (dialogue) and output (summary) from the dataset
-examples = [
-    {
-        "input": dialogue['dialogue'],  # Assuming 'dialogue' field contains the conversation text
-        "output": dialogue['summary']   # Assuming 'summary' field contains the summary
-    }
-    for dialogue in dataset
-]
+examples = load_examples()
 
-# Create FewShotChatMessagePromptTemplate without prefix and suffix
-example_prompt = FewShotChatMessagePromptTemplate(
+# 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,
-    input_variables=["input"]  # only the required fields
+    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.title("📝 Text Summarizer using Few-Shot Prompt")
-
-input_text = st.text_area("Enter the text you want to summarize:")
-
-if st.button("Summarize"):
-    if input_text.strip():
-        # Format the prompt
-        formatted_message = example_prompt.format(input=input_text)
-
-        with st.expander("🔍 Prompt Preview"):
-            st.markdown(f"**Formatted Prompt:** {formatted_message}")
-
-        # Load the model from Hugging Face (replace with your choice of model)
-        model = HuggingFaceHub(
-            repo_id="google/pegasus-xsum",  # You can replace with any model available in Hugging Face
-            model_kwargs={"temperature": 0.7}
-        )
-
-        # Generate the summary
-        summary = model(formatted_message)
-        st.success("✅ Summary:")
-        st.write(summary)
-    else:
-        st.warning("Please enter some text!")
+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)