Update app.py

app.py

@@ -1,54 +1,50 @@
 import streamlit as st
 from langchain.prompts import FewShotChatMessagePromptTemplate
-from langchain.prompts.example_selector import LengthBasedExampleSelector
 from langchain.llms import HuggingFaceHub
 from datasets import load_dataset
 
-# Load dataset (you can use any summarization dataset here)
-@st.cache_data
-def load_examples():
-    # Using 'knkarthick/dialogsum' as an example dataset
-    dataset = load_dataset("knkarthick/dialogsum", split="train[:5]")  # Load a subset for testing
-    examples = []
-    for example in dataset:
-        examples.append({
-            "input": example["dialogue"],
-            "output": example["summary"]
-        })
-    return examples
-
-# Load few-shot examples from the dataset
-examples = load_examples()
+# Load the dataset
+dataset = load_dataset("knkarthick/dialogsum", split="train[:1%]")  # Load a small subset of the dataset for testing
+
+# 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
+]
 
 # Create FewShotChatMessagePromptTemplate
-example_prompt = FewShotChatMessagePromptTemplate.from_examples(
+example_prompt = FewShotChatMessagePromptTemplate(
     examples=examples,
-    example_selector=LengthBasedExampleSelector(examples=examples, max_length=1000),
     input_variables=["input"],
-    prefix="You are a helpful assistant that summarizes dialogues. Examples:",
-    suffix="Now summarize this:\n{input}"
+    prefix="You are a helpful summarizer. Here are a few examples:",
+    suffix="Now summarize this: {input}"
 )
 
-# Set up Hugging Face model (you can replace it with any other available model)
-llm = HuggingFaceHub(repo_id="t5-small", task="summarization")
-
-# Streamlit UI setup
-st.title("📝 Dialogue Summarizer using Few-Shot Prompt + T5")
+# Streamlit UI
+st.title("📝 Text Summarizer using Few-Shot Prompt")
 
-input_text = st.text_area("📝 Paste your conversation here:")
+input_text = st.text_area("Enter the text you want to summarize:")
 
-if st.button("Generate Summary"):
+if st.button("Summarize"):
     if input_text.strip():
-        # Create the prompt using FewShotChatMessagePromptTemplate
-        messages = example_prompt.format_messages(input=input_text)
+        # Format the prompt
+        formatted_message = example_prompt.format(input=input_text)
+
+        with st.expander("🔍 Prompt Preview"):
+            st.markdown(f"**Formatted Prompt:** {formatted_message}")
 
-        with st.expander("📋 Generated Prompt"):
-            for msg in messages:
-                st.markdown(f"**{msg.type.upper()}**:\n```\n{msg.content}\n```")
+        # 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}
+        )
 
-        # Get the summary using the Hugging Face model
-        response = llm(messages[0].content)
+        # Generate the summary
+        summary = model(formatted_message)
         st.success("✅ Summary:")
-        st.write(response['text'])
+        st.write(summary)
     else:
-        st.warning("Please enter some text.")
+        st.warning("Please enter some text!")