Update app.py

app.py

@@ -1,41 +1,52 @@
 import streamlit as st
+from datasets import load_dataset
+from transformers import pipeline
 from langchain.prompts.few_shot import FewShotChatMessagePromptTemplate
 from langchain.prompts.example_selector import LengthBasedExampleSelector
 
-# Example dialogues + summaries
-examples = [
-    {
-        "input": "Doctor: What symptoms are you experiencing?\nPatient: I have a sore throat and runny nose.",
-        "output": "The patient reports sore throat and runny nose."
-    },
-    {
-        "input": "Patient: I've been feeling dizzy since morning.\nDoctor: Did you eat anything unusual?\nPatient: No, just normal food.",
-        "output": "The patient has dizziness but no unusual food intake."
-    }
-]
-
-# Create Few-Shot Prompt Template
-prompt = FewShotChatMessagePromptTemplate.from_examples(
+# 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 summarization model
+summarizer = pipeline("summarization", model="t5-small")
+
+# Few-shot prompt template
+example_prompt = FewShotChatMessagePromptTemplate.from_examples(
     examples=examples,
-    example_selector=LengthBasedExampleSelector(examples=examples, max_length=100),
+    example_selector=LengthBasedExampleSelector(examples=examples, max_length=1000),
     input_variables=["input"],
-    prefix="You are a medical assistant that summarizes doctor-patient conversations. Examples:",
-    suffix="Now summarize this conversation:\n{input}"
+    prefix="You are a helpful assistant that summarizes dialogues. Examples:",
+    suffix="Now summarize this:\n{input}"
 )
 
 # Streamlit UI
-st.title("💬 Few-Shot Summarizer (No Model)")
-
-input_text = st.text_area("Paste your medical conversation here:")
+st.title("💬 Dialogue Summarizer using Few-Shot Prompt + T5")
+input_text = st.text_area("📝 Paste your conversation:")
 
-if st.button("Generate Prompt"):
+if st.button("Generate Summary"):
     if input_text.strip():
-        messages = prompt.format_messages(input=input_text)
+        # Create prompt
+        messages = example_prompt.format_messages(input=input_text)
 
-        st.subheader("📋 Prompt Generated:")
-        for msg in messages:
-            st.markdown(f"**{msg.type.upper()}**:\n```\n{msg.content}\n```")
+        with st.expander("📋 Generated Prompt"):
+            for msg in messages:
+                st.markdown(f"**{msg.type.upper()}**:\n```\n{msg.content}\n```")
 
-        st.info("⚠️ This is just a prompt, not an actual summary.")
+        # Generate summary using model
+        full_prompt = "\n".join([m.content for m in messages if m.type == "human"])
+        summary = summarizer("summarize: " + input_text, max_length=80, min_length=15, do_sample=False)[0]['summary_text']
+        st.success("✅ Summary:")
+        st.write(summary)
     else:
-        st.warning("Please enter some conversation text.")
+        st.warning("Please enter some text.")