Update app.py

app.py

@@ -1,9 +1,9 @@
 import os
 import streamlit as st
 import torch
-from langchain.chains import LLMChain
 from langchain.prompts import ChatPromptTemplate
-from langchain_community.llms import HuggingFaceHub
+from langchain.schema.runnable import RunnableSequence
+from langchain_huggingface import HuggingFaceEndpoint
 
 def create_prompt(name: str, persona_style: str):
     """Create the chat prompt template."""
@@ -32,7 +32,7 @@ def create_prompt(name: str, persona_style: str):
     """
     return ChatPromptTemplate.from_template(prompt_template_str)
 
-def simulate_conversation(chain: LLMChain, turns: int = 15, max_history_rounds=3):
+def simulate_conversation(chain: RunnableSequence, 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 +57,11 @@ def simulate_conversation(chain: LLMChain, turns: int = 15, max_history_rounds=3
         for i in range(turns):
             human_input = human_messages[i % len(human_messages)]
             
-            # Build truncated chat_history for prompt
-            # Keep only the last max_history_rounds * 2 lines (Human + AI pairs)
+            # Keep only last max_history_rounds * 2 lines
             truncated_history_lines = chat_history_list[-(max_history_rounds*2):]
             truncated_history = "\n".join(truncated_history_lines)
 
-            response = chain.run(chat_history=truncated_history, input=human_input)
+            response = chain.invoke({"chat_history": truncated_history, "input": human_input})
             # Update chat history
             chat_history_list.append(f"Human: {human_input}")
             chat_history_list.append(f"AI: {response}")
@@ -73,11 +72,11 @@ def simulate_conversation(chain: LLMChain, turns: int = 15, max_history_rounds=3
         st.error(f"Error during conversation simulation: {e}")
         return None
 
-def summarize_conversation(chain: LLMChain, conversation: str):
+def summarize_conversation(chain: RunnableSequence, 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.run(chat_history="", input=summary_prompt)
+        response = chain.invoke({"chat_history": "", "input": summary_prompt})
         return response.strip()
     except Exception as e:
         st.error(f"Error summarizing conversation: {e}")
@@ -86,7 +85,6 @@ def summarize_conversation(chain: LLMChain, conversation: str):
 def main():
     st.title("LLM Conversation Simulation")
 
-    # Model selection
     model_names = [
         "meta-llama/Llama-3.3-70B-Instruct",
         "meta-llama/Llama-3.1-405B-Instruct",
@@ -94,38 +92,44 @@ def main():
     ]
     selected_model = st.selectbox("Select a model:", model_names)
 
-    # Persona Inputs
     name = st.text_input("Enter the persona's name:", value="Alex")
     persona_style = st.text_area("Enter the persona style characteristics:", 
                                  value="friendly, curious, and a bit sarcastic")
 
     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
+            
+            endpoint_url = f"https://api-inference.huggingface.co/models/{selected_model}"
+            headers = {"Authorization": f"Bearer {hf_token}"}
+            
             try:
-                # Use HuggingFaceHub as LLM
-                # Make sure you have a valid HUGGINGFACEHUB_API_TOKEN set
-                llm = HuggingFaceHub(
-                    repo_id=selected_model,
+                llm = HuggingFaceEndpoint(
+                    endpoint_url=endpoint_url,
+                    task="text-generation",
+                    headers=headers,
                     model_kwargs={
                         "temperature": 0.7,
                         "max_new_tokens": 512
                     }
                 )
             except Exception as e:
-                st.error(f"Error initializing model from Hugging Face Hub: {e}")
+                st.error(f"Error initializing HuggingFaceEndpoint: {e}")
                 return
 
-            # Create our prompt template chain
             prompt = create_prompt(name, persona_style)
-            chain = LLMChain(llm=llm, prompt=prompt)
+            # prompt and llm are both Runnables, chain them together
+            chain = RunnableSequence([prompt, llm])
 
-            # Simulate conversation
             conversation = simulate_conversation(chain, turns=15, max_history_rounds=3)
             if conversation:
                 st.subheader("Conversation:")
                 st.text(conversation)
 
-                # Summarize conversation
                 st.subheader("Summary:")
                 summary = summarize_conversation(chain, conversation)
                 st.write(summary)