Update app.py

app.py

@@ -27,18 +27,10 @@ few_shot_prompt = FewShotChatMessagePromptTemplate(
 # 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
+# Streamlit UI setup
 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.")
@@ -46,12 +38,21 @@ 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)
+    # Prepare messages for the final prompt (include few-shot examples here directly)
+    formatted_prompt = few_shot_prompt.format_messages(dialogue=user_input)
     
-    # Output
+    # Add formatted examples to the final prompt
+    formatted_prompt = final_prompt.format_messages(dialogue=user_input)
+    formatted_prompt = formatted_prompt + formatted_prompt
+
+    # Get response from model
+    llm = HuggingFaceEndpoint(
+        repo_id="google/pegasus-xsum",
+        task="text2text-generation",
+        model_kwargs={"temperature": 0.3, "max_new_tokens": 128}
+    )
+    response = llm(formatted_prompt)
+
+    # Output the summary
     st.subheader("📌 Summary:")
     st.write(response)