Update app.py

app.py

@@ -2,8 +2,9 @@ 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 dataset from HuggingFace
+# Load few-shot examples from dialogsum
 @st.cache_data
 def load_examples(n=3):
     dataset = load_dataset("knkarthick/dialogsum", split="train[:20]")
@@ -11,22 +12,26 @@ def load_examples(n=3):
 
 examples = load_examples()
 
-# Format examples
+# Template for each example
 example_prompt = ChatPromptTemplate.from_messages([
     ("human", "Summarize the following dialog:\n\n{dialogue}"),
     ("ai", "{summary}")
 ])
 
-# Few-shot setup
+# Few-shot prompt template (no prefix/suffix here)
 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."
+    examples=examples
 )
 
-# Load HF summarizer model (Pegasus)
+# 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.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",
@@ -41,7 +46,12 @@ 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)
+    # Format messages
+    messages = final_prompt.format_messages(dialogue=user_input)
+    
+    # Get response
     response = llm(messages)
+    
+    # Output
     st.subheader("📌 Summary:")
     st.write(response)