Update app.py

app.py

@@ -1,8 +1,10 @@
 import os
 import streamlit as st
 import torch
+from langchain.chains import LLMChain
 from langchain.prompts import ChatPromptTemplate
-from langchain.schema.runnable import RunnableSequence
+
+# Use the new package for HuggingFaceEndpoint
 from langchain_huggingface import HuggingFaceEndpoint
 
 def create_prompt(name: str, persona_style: str):
@@ -32,7 +34,7 @@ def create_prompt(name: str, persona_style: str):
     """
     return ChatPromptTemplate.from_template(prompt_template_str)
 
-def simulate_conversation(chain: RunnableSequence, turns: int = 15, max_history_rounds=3):
+def simulate_conversation(chain: LLMChain, turns: int = 15, max_history_rounds=3):
     """Simulate a conversation for a given number of turns, limiting chat history."""
     chat_history_list = []
     human_messages = [
@@ -57,12 +59,11 @@ def simulate_conversation(chain: RunnableSequence, turns: int = 15, max_history_
         for i in range(turns):
             human_input = human_messages[i % len(human_messages)]
             
-            # Keep only last max_history_rounds * 2 lines
+            # Build truncated chat_history for prompt
             truncated_history_lines = chat_history_list[-(max_history_rounds*2):]
             truncated_history = "\n".join(truncated_history_lines)
 
-            response = chain.invoke({"chat_history": truncated_history, "input": human_input})
-            # Update chat history
+            response = chain.run(chat_history=truncated_history, input=human_input)
             chat_history_list.append(f"Human: {human_input}")
             chat_history_list.append(f"AI: {response}")
 
@@ -72,11 +73,11 @@ def simulate_conversation(chain: RunnableSequence, turns: int = 15, max_history_
         st.error(f"Error during conversation simulation: {e}")
         return None
 
-def summarize_conversation(chain: RunnableSequence, conversation: str):
+def summarize_conversation(chain: LLMChain, conversation: str):
     """Use the LLM to summarize the completed conversation."""
     summary_prompt = f"Summarize the following conversation in a few short sentences highlighting the main points, tone, and conclusion:\n\n{conversation}\nSummary:"
     try:
-        response = chain.invoke({"chat_history": "", "input": summary_prompt})
+        response = chain.run(chat_history="", input=summary_prompt)
         return response.strip()
     except Exception as e:
         st.error(f"Error summarizing conversation: {e}")
@@ -98,32 +99,25 @@ def main():
 
     if st.button("Start Conversation Simulation"):
         with st.spinner("Starting simulation..."):
-            # Build headers with your Hugging Face token
-            hf_token = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
-            if not hf_token:
-                st.error("HUGGINGFACEHUB_API_TOKEN not found. Please set the token.")
-                return
-            
+            # Construct the endpoint URL for the selected model
             endpoint_url = f"https://api-inference.huggingface.co/models/{selected_model}"
-            headers = {"Authorization": f"Bearer {hf_token}"}
-            
+
             try:
+                # Use HuggingFaceEndpoint instead of HuggingFaceHub
+                # Specify temperature and max_new_tokens as top-level arguments
                 llm = HuggingFaceEndpoint(
                     endpoint_url=endpoint_url,
+                    huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN"),
                     task="text-generation",
-                    headers=headers,
-                    model_kwargs={
-                        "temperature": 0.7,
-                        "max_new_tokens": 512
-                    }
+                    temperature=0.7,
+                    max_new_tokens=512
                 )
             except Exception as e:
                 st.error(f"Error initializing HuggingFaceEndpoint: {e}")
                 return
 
             prompt = create_prompt(name, persona_style)
-            # prompt and llm are both Runnables, chain them together
-            chain = RunnableSequence([prompt, llm])
+            chain = LLMChain(llm=llm, prompt=prompt)
 
             conversation = simulate_conversation(chain, turns=15, max_history_rounds=3)
             if conversation: