Update app.py

app.py

@@ -1,12 +1,13 @@
 import streamlit as st
-from datasets import load_dataset
+from langchain.prompts import FewShotChatMessagePromptTemplate
+from langchain.prompts.example_selector import LengthBasedExampleSelector
+from langchain_huggingface import HuggingFaceEndpoint, HuggingFacePipeline
 from langchain.chains import LLMChain
-from langchain.llms import HuggingFaceHub
 from langchain.prompts import PromptTemplate
-from langchain.prompts.few_shot import FewShotChatMessagePromptTemplate
-from langchain.prompts.example_selector import LengthBasedExampleSelector
+from datasets import load_dataset
+from transformers import pipeline
 
-# Load dataset (small subset)
+# Load dataset (using knkarthick/dialogsum as an example)
 @st.cache_data
 def load_examples():
     dataset = load_dataset("knkarthick/dialogsum", split="train[:5]")  # Take only 5 for speed
@@ -20,10 +21,12 @@ def load_examples():
 
 examples = load_examples()
 
-# Set up the HuggingFaceHub model (T5)
-llm = HuggingFaceHub(repo_id="google/pegasus-xsum", model_kwargs={"temperature": 0.7})
+# Load the Hugging Face model
+hf_endpoint = HuggingFaceEndpoint(
+    endpoint_url="https://api-inference.huggingface.co/models/t5-small"  # or any model you like
+)
 
-# Few-shot prompt template
+# Create FewShotChatMessagePromptTemplate
 example_prompt = FewShotChatMessagePromptTemplate.from_examples(
     examples=examples,
     example_selector=LengthBasedExampleSelector(examples=examples, max_length=1000),
@@ -34,24 +37,24 @@ example_prompt = FewShotChatMessagePromptTemplate.from_examples(
 
 # Streamlit UI
 st.title("💬 Dialogue Summarizer using Few-Shot Prompt + T5 (via Langchain)")
+
 input_text = st.text_area("📝 Paste your conversation:")
 
 if st.button("Generate Summary"):
     if input_text.strip():
-        # Create prompt using FewShotChatMessagePromptTemplate
+        # Create the prompt using FewShotChatMessagePromptTemplate
         messages = example_prompt.format_messages(input=input_text)
 
         with st.expander("📋 Generated Prompt"):
             for msg in messages:
                 st.markdown(f"**{msg.type.upper()}**:\n```\n{msg.content}\n```")
 
-        # Create the prompt chain
-        prompt_template = PromptTemplate(input_variables=["input"], template="{input}")
-        chain = LLMChain(llm=llm, prompt=prompt_template)
+        # Set up HuggingFacePipeline with the model endpoint
+        hf_pipeline = HuggingFacePipeline(pipeline="summarization", model=hf_endpoint)
 
-        # Get the summary from the model
-        summary = chain.run(input_text)
+        # Generate summary
+        summary = hf_pipeline(messages[0].content)
         st.success("✅ Summary:")
-        st.write(summary)
+        st.write(summary[0]['summary_text'])
     else:
         st.warning("Please enter some text.")