Update app.py

app.py

@@ -1,208 +1,16 @@
-# import os
-# import streamlit as st
-# import torch
-# from langchain.chains import LLMChain
-# from langchain.prompts import ChatPromptTemplate
-# from langchain_huggingface import HuggingFaceEndpoint
-
-# def create_conversation_prompt(name1: str, name2: str, persona_style: str):
-#     """
-#     Create a prompt that instructs the model to produce exactly 15 messages
-#     of conversation, alternating between name1 and name2, starting with name1.
-
-#     We will be very explicit and not allow any formatting except the required lines.
-#     """
-#     prompt_template_str = f"""
-#     You are simulating a conversation of exactly 15 messages between two people: {name1} and {name2}.
-#     {name1} speaks first (message 1), then {name2} (message 2), then {name1} (message 3), and so forth,
-#     alternating until all 15 messages are complete. The 15th message is by {name1}.
-
-#     Requirements:
-#     - Output exactly 15 lines, no more, no less.
-#     - Each line must be a single message in the format:
-#       {name1}: <message> or {name2}: <message>
-#     - Do not add any headings, numbers, sample outputs, or explanations.
-#     - Do not mention code, programming, or instructions.
-#     - Each message should be 1-2 short sentences, friendly, natural, reflecting the style: {persona_style}.
-#     - Use everyday language, can ask questions, show opinions.
-#     - Use emojis sparingly if it fits the style (no more than 1-2 total).
-#     - No repeated lines, each message should logically follow from the previous one.
-#     - Do not produce anything after the 15th message. No extra lines or text.
-
-#     Produce all 15 messages now:
-#     """
-#     return ChatPromptTemplate.from_template(prompt_template_str)
-
-# def create_summary_prompt(name1: str, name2: str, conversation: str):
-#     """Prompt for generating a title and summary."""
-#     summary_prompt_str = f"""
-#     Below is a completed 15-message conversation between {name1} and {name2}:
-
-#     {conversation}
-
-#     Please provide:
-#     Title: <A short descriptive title of the conversation>
-#     Summary: <A few short sentences highlighting the main points, tone, and conclusion>
-
-#     Do not continue the conversation, do not repeat it, and do not add extra formatting beyond the two lines:
-#     - One line starting with "Title:"
-#     - One line starting with "Summary:"
-#     """
-#     return ChatPromptTemplate.from_template(summary_prompt_str)
-
-# def main():
-#     st.title("LLM Conversation Simulation")
-
-#     model_names = [
-#         "meta-llama/Llama-3.3-70B-Instruct",
-#         "meta-llama/Llama-3.1-405B-Instruct",
-#         "Qwen/Qwen2.5-72B-Instruct",
-#         "deepseek-ai/DeepSeek-V3",
-#         "deepseek-ai/DeepSeek-V2.5"
-        
-#     ]
-#     selected_model = st.selectbox("Select a model:", model_names)
-
-#     name1 = st.text_input("Enter the first user's name:", value="Alice")
-#     name2 = st.text_input("Enter the second user's name:", value="Bob")
-#     persona_style = st.text_area("Enter the persona style characteristics:", 
-#                                  value="friendly, curious, and a bit sarcastic")
-
-#     if st.button("Start Conversation Simulation"):
-#         st.write("**Loading model...**")
-#         print("Loading model...")
-
-#         with st.spinner("Starting simulation..."):
-#             endpoint_url = f"https://api-inference.huggingface.co/models/{selected_model}"
-
-#             try:
-#                 llm = HuggingFaceEndpoint(
-#                     endpoint_url=endpoint_url,
-#                     huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN"),
-#                     task="text-generation",
-#                     temperature=0.7,
-#                     max_new_tokens=512
-#                 )
-#                 st.write("**Model loaded successfully!**")
-#                 print("Model loaded successfully!")
-#             except Exception as e:
-#                 st.error(f"Error initializing HuggingFaceEndpoint: {e}")
-#                 print(f"Error initializing HuggingFaceEndpoint: {e}")
-#                 return
-
-#             conversation_prompt = create_conversation_prompt(name1, name2, persona_style)
-#             conversation_chain = LLMChain(llm=llm, prompt=conversation_prompt)
-
-#             st.write("**Generating the full 15-message conversation...**")
-#             print("Generating the full 15-message conversation...")
-
-#             try:
-#                 # Generate all 15 messages in one go
-#                 conversation = conversation_chain.run(chat_history="", input="").strip()
-
-#                 st.subheader("Final Conversation:")
-#                 st.text(conversation)
-#                 print("Conversation Generation Complete.\n")
-#                 print("Full Conversation:\n", conversation)
-
-#                 # Summarize the conversation
-#                 summary_prompt = create_summary_prompt(name1, name2, conversation)
-#                 summary_chain = LLMChain(llm=llm, prompt=summary_prompt)
-
-#                 st.subheader("Summary and Title:")
-#                 st.write("**Summarizing the conversation...**")
-#                 print("Summarizing the conversation...")
-
-#                 summary = summary_chain.run(chat_history="", input="")
-#                 st.write(summary)
-#                 print("Summary:\n", summary)
-
-#             except Exception as e:
-#                 st.error(f"Error generating conversation: {e}")
-#                 print(f"Error generating conversation: {e}")
-
-# if __name__ == "__main__":
-#     main()
-
-
 import os
 import streamlit as st
-import google.cloud.aiplatform as aiplatform
+import torch
 from langchain.chains import LLMChain
 from langchain.prompts import ChatPromptTemplate
-from langchain.llms.base import LLM
-from pydantic import BaseModel
-from typing import Optional, List, Mapping, Any
+from langchain_huggingface import HuggingFaceEndpoint
 
-###############################################################################
-# 1. Create a Custom LLM class for LangChain to call your Vertex AI endpoint.
-###############################################################################
-class VertexAICustomModel(LLM, BaseModel):
-    project_id: str
-    location: str
-    endpoint_id: str
-    temperature: float = 0.7
-    max_new_tokens: int = 512
-    
-    @property
-    def _llm_type(self) -> str:
-        return "vertex_ai_custom"
-
-    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
-        # Initialize Vertex AI with your project/region
-        aiplatform.init(project=self.project_id, location=self.location)
-        endpoint = aiplatform.Endpoint(
-            endpoint_name=f"projects/{self.project_id}/locations/{self.location}/endpoints/{self.endpoint_id}"
-        )
-        
-        # Construct the instance for prediction. 
-        # NOTE: Adjust 'prompt', 'temperature', etc. if your model expects different parameters.
-        instance = {
-            "prompt": prompt,
-            "temperature": self.temperature,
-            "max_new_tokens": self.max_new_tokens
-        }
-
-        # Call the endpoint
-        response = endpoint.predict(instances=[instance])
-        
-        # Extract the text from the response.
-        # This will vary depending on how your model returns predictions.
-        # A common approach is response.predictions[0]["generated_text"], 
-        # but confirm your model's actual JSON structure.
-        predictions = response.predictions
-        if not predictions or "generated_text" not in predictions[0]:
-            raise ValueError(
-                f"Unexpected response structure from Vertex AI endpoint: {response}"
-            )
-        
-        text = predictions[0]["generated_text"]
-
-        # Optionally apply 'stop' tokens
-        if stop:
-            for s in stop:
-                if s in text:
-                    text = text.split(s)[0]
-        return text
-
-    @property
-    def _identifying_params(self) -> Mapping[str, Any]:
-        """Return any identifying parameters of this LLM."""
-        return {
-            "endpoint_id": self.endpoint_id,
-            "project_id": self.project_id,
-            "location": self.location,
-            "temperature": self.temperature,
-            "max_new_tokens": self.max_new_tokens,
-        }
-
-###############################################################################
-# 2. Create your conversation and summary prompt templates (unchanged).
-###############################################################################
 def create_conversation_prompt(name1: str, name2: str, persona_style: str):
     """
     Create a prompt that instructs the model to produce exactly 15 messages
     of conversation, alternating between name1 and name2, starting with name1.
+
+    We will be very explicit and not allow any formatting except the required lines.
     """
     prompt_template_str = f"""
     You are simulating a conversation of exactly 15 messages between two people: {name1} and {name2}.
@@ -229,77 +37,89 @@ def create_summary_prompt(name1: str, name2: str, conversation: str):
     """Prompt for generating a title and summary."""
     summary_prompt_str = f"""
     Below is a completed 15-message conversation between {name1} and {name2}:
+
     {conversation}
+
     Please provide:
     Title: <A short descriptive title of the conversation>
     Summary: <A few short sentences highlighting the main points, tone, and conclusion>
+
     Do not continue the conversation, do not repeat it, and do not add extra formatting beyond the two lines:
     - One line starting with "Title:"
     - One line starting with "Summary:"
     """
     return ChatPromptTemplate.from_template(summary_prompt_str)
 
-###############################################################################
-# 3. Main Streamlit app with Vertex AI usage.
-###############################################################################
 def main():
-    st.title("LLM Conversation Simulation (GCP Vertex AI)")
-
-    # We can remove model selection if we are always using your deployed model:
-    # st.selectbox(... ) # => Removed
-
-    # Hardcode or load your Vertex AI endpoint details here
-    project_id = "282802344966"
-    location = "us-west1"
-    endpoint_id = "1106913540054188032"
+    st.title("LLM Conversation Simulation")
+
+    model_names = [
+        "meta-llama/Llama-3.3-70B-Instruct",
+        "meta-llama/Llama-3.1-405B-Instruct",
+        "Qwen/Qwen2.5-72B-Instruct",
+        "deepseek-ai/DeepSeek-V3",
+        "deepseek-ai/DeepSeek-V2.5"
+        
+    ]
+    selected_model = st.selectbox("Select a model:", model_names)
 
-    # Input fields for conversation
     name1 = st.text_input("Enter the first user's name:", value="Alice")
     name2 = st.text_input("Enter the second user's name:", value="Bob")
-    persona_style = st.text_area("Enter the persona style characteristics:",
+    persona_style = st.text_area("Enter the persona style characteristics:", 
                                  value="friendly, curious, and a bit sarcastic")
 
     if st.button("Start Conversation Simulation"):
-        st.write("**Initializing Vertex AI endpoint...**")
-        st.spinner("Starting simulation...")
-
-        # Create your custom LLM that calls Vertex AI
-        llm = VertexAICustomModel(
-            project_id=project_id,
-            location=location,
-            endpoint_id=endpoint_id,
-            temperature=0.7,
-            max_new_tokens=512
-        )
-
-        st.write("**Vertex AI endpoint loaded successfully!**")
-
-        # Build the conversation chain
-        conversation_prompt = create_conversation_prompt(name1, name2, persona_style)
-        conversation_chain = LLMChain(llm=llm, prompt=conversation_prompt)
-
-        st.write("**Generating the full 15-message conversation...**")
-
-        try:
-            # Generate all 15 messages in one go
-            conversation = conversation_chain.run(chat_history="", input="").strip()
-
-            st.subheader("Final Conversation:")
-            st.text(conversation)
-
-            # Summarize the conversation
-            summary_prompt = create_summary_prompt(name1, name2, conversation)
-            summary_chain = LLMChain(llm=llm, prompt=summary_prompt)
-
-            st.subheader("Summary and Title:")
-            st.write("**Summarizing the conversation...**")
-
-            summary = summary_chain.run(chat_history="", input="")
-            st.write(summary)
-
-        except Exception as e:
-            st.error(f"Error generating conversation: {e}")
-            print(f"Error generating conversation: {e}")
+        st.write("**Loading model...**")
+        print("Loading model...")
+
+        with st.spinner("Starting simulation..."):
+            endpoint_url = f"https://api-inference.huggingface.co/models/{selected_model}"
+
+            try:
+                llm = HuggingFaceEndpoint(
+                    endpoint_url=endpoint_url,
+                    huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN"),
+                    task="text-generation",
+                    temperature=0.7,
+                    max_new_tokens=512
+                )
+                st.write("**Model loaded successfully!**")
+                print("Model loaded successfully!")
+            except Exception as e:
+                st.error(f"Error initializing HuggingFaceEndpoint: {e}")
+                print(f"Error initializing HuggingFaceEndpoint: {e}")
+                return
+
+            conversation_prompt = create_conversation_prompt(name1, name2, persona_style)
+            conversation_chain = LLMChain(llm=llm, prompt=conversation_prompt)
+
+            st.write("**Generating the full 15-message conversation...**")
+            print("Generating the full 15-message conversation...")
+
+            try:
+                # Generate all 15 messages in one go
+                conversation = conversation_chain.run(chat_history="", input="").strip()
+
+                st.subheader("Final Conversation:")
+                st.text(conversation)
+                print("Conversation Generation Complete.\n")
+                print("Full Conversation:\n", conversation)
+
+                # Summarize the conversation
+                summary_prompt = create_summary_prompt(name1, name2, conversation)
+                summary_chain = LLMChain(llm=llm, prompt=summary_prompt)
+
+                st.subheader("Summary and Title:")
+                st.write("**Summarizing the conversation...**")
+                print("Summarizing the conversation...")
+
+                summary = summary_chain.run(chat_history="", input="")
+                st.write(summary)
+                print("Summary:\n", summary)
+
+            except Exception as e:
+                st.error(f"Error generating conversation: {e}")
+                print(f"Error generating conversation: {e}")
 
 if __name__ == "__main__":
     main()