Update app.py

app.py

@@ -2,48 +2,10 @@ import streamlit as st
 import torch
 from langchain.chains import LLMChain
 from langchain.prompts import ChatPromptTemplate
-from langchain_community.llms import HuggingFacePipeline  # Updated import
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-
-# ---------------------------------------
-# Utility functions
-# ---------------------------------------
-
-def load_model_pipeline(model_name: str):
-    """Load a HuggingFace pipeline for a given model name.
-    
-    For large models, ensure you have the model downloaded and configured.
-    Adjust device_map, load_in_8bit, etc., as needed for your hardware.
-    """
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name, 
-            device_map="auto",  # Adjust as needed (e.g. "cpu", "cuda:0")
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            trust_remote_code=True
-        )
-        # Create pipeline
-        pipe = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            max_length=512,
-            temperature=0.7,
-            top_p=0.9,
-            repetition_penalty=1.1,
-            pad_token_id=tokenizer.eos_token_id,
-        )
-        return pipe
-    except Exception as e:
-        st.error(f"Error loading model {model_name}: {e}")
-        return None
-
+from langchain_community.llms import HuggingFaceHub  # Using HuggingFaceHub now
 
 def create_prompt(name: str, persona_style: str):
     """Create the chat prompt template as described."""
-    # The prompt template from the user prompt
-    # We combine both snippet versions into one final template.
     prompt_template_str = f"""
     In this interaction, your role is to embody the personality of a human named {name}, with characteristics defined as {persona_style}. Engage in a natural, human-like conversation, responding to each query or comment as if you were {name} with the specific traits and style of {persona_style}. Remember, your responses should reflect a human perspective, complete with emotions, opinions, and personal insights that align with {persona_style}. ////
     //////////
@@ -69,19 +31,9 @@ def create_prompt(name: str, persona_style: str):
     """
     return ChatPromptTemplate.from_template(prompt_template_str)
 
-
 def simulate_conversation(chain: LLMChain, turns: int = 15):
-    """Simulate a conversation for a given number of turns.
-    
-    Here, we'll pretend the "Human" says something generic each time,
-    and we get the AI's response. We store and update the chat history.
-    
-    After 15 responses from the AI, we return the full conversation.
-    """
-
+    """Simulate a conversation for a given number of turns."""
     chat_history = ""
-    # We will simulate the human input as a rotating set of simple messages 
-    # or just a single repeated message to show the flow.
     human_messages = [
         "Hey, what's up?",
         "That's interesting, tell me more!",
@@ -103,19 +55,15 @@ def simulate_conversation(chain: LLMChain, turns: int = 15):
     try:
         for i in range(turns):
             human_input = human_messages[i % len(human_messages)]
-            # Generate AI response
             response = chain.run(chat_history=chat_history, input=human_input)
-            # Update the chat history
             chat_history += f"Human: {human_input}\nAI: {response}\n"
         return chat_history
     except Exception as e:
         st.error(f"Error during conversation simulation: {e}")
         return None
 
-
 def summarize_conversation(chain: LLMChain, conversation: str):
     """Use the LLM to summarize the completed conversation."""
-    # We'll provide a simple prompt for summary:
     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)
@@ -124,14 +72,9 @@ def summarize_conversation(chain: LLMChain, conversation: str):
         st.error(f"Error summarizing conversation: {e}")
         return "No summary available due to error."
 
-
-# ---------------------------------------
-# Streamlit App
-# ---------------------------------------
 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",
@@ -139,34 +82,39 @@ 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")
 
-    # Button to start simulation
     if st.button("Start Conversation Simulation"):
-        with st.spinner("Loading model and starting simulation..."):
-            pipe = load_model_pipeline(selected_model)
-            if pipe is not None:
-                # Create a ChatModel from the pipeline
-                llm = HuggingFacePipeline(pipeline=pipe)
-                
-                # Create our prompt template chain
-                prompt = create_prompt(name, persona_style)
-                chain = LLMChain(llm=llm, prompt=prompt)
-
-                # Simulate conversation
-                conversation = simulate_conversation(chain, turns=15)
-                if conversation:
-                    st.subheader("Conversation:")
-                    st.text(conversation)
-                    
-                    # Summarize conversation
-                    st.subheader("Summary:")
-                    summary = summarize_conversation(chain, conversation)
-                    st.write(summary)
-
+        with st.spinner("Starting simulation..."):
+            try:
+                # Using HuggingFaceHub for remote model inference
+                llm = HuggingFaceHub(
+                    repo_id=selected_model,
+                    model_kwargs={
+                        "temperature": 0.7,
+                        "max_new_tokens": 512
+                    }
+                )
+            except Exception as e:
+                st.error(f"Error initializing model from Hugging Face Hub: {e}")
+                return
+
+            # Create our prompt template chain
+            prompt = create_prompt(name, persona_style)
+            chain = LLMChain(llm=llm, prompt=prompt)
+
+            # Simulate conversation
+            conversation = simulate_conversation(chain, turns=15)
+            if conversation:
+                st.subheader("Conversation:")
+                st.text(conversation)
+
+                # Summarize conversation
+                st.subheader("Summary:")
+                summary = summarize_conversation(chain, conversation)
+                st.write(summary)
 
 if __name__ == "__main__":
     main()