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 for generating the entire 15-message conversation.""" prompt_template_str = f""" You are to simulate a conversation of exactly 15 messages total between two people: {name1} and {name2}. The conversation should reflect the style: {persona_style}. {name1} speaks first (message 1), {name2} responds (message 2), then {name1} (message 3), and so forth, until 15 messages are complete (the 15th message by {name1}). Rules: - Each message is formatted as: {name1}: or {name2}: - Each message: 1-2 short sentences, friendly, natural. - Use everyday language, can ask questions, show opinions. - Use emojis sparingly if it fits the style. - Do not repeat the same line. - Produce all 15 messages now and do not continue beyond the 15th message. """ return ChatPromptTemplate.from_template(prompt_template_str) def create_summary_prompt(name1: str, name2: str, conversation: str): """Create a prompt specifically for generating a title and summary of the conversation.""" # Here we explicitly create a new prompt template for the summary. summary_prompt_str = f""" The following is a completed conversation between {name1} and {name2}: {conversation} Please provide: Title: Summary: Do not continue the conversation, just provide title and summary. """ return ChatPromptTemplate.from_template(summary_prompt_str) def main(): st.title("LLM Conversation Simulation") model_names = [ "meta-llama/Llama-3.3-70B-Instruct", "mistralai/Mistral-7B-v0.1", "lmsys/vicuna-13b-v1.5", "tiiuae/falcon-180B", "EleutherAI/gpt-neox-20b", "dice-research/lola_v1" ] 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 # Create a chain for the conversation generation 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 # Here we send the prompt for the conversation to the LLM conversation = conversation_chain.run(chat_history="", input="Produce the full conversation now.") conversation = conversation.strip() st.subheader("Final Conversation:") st.text(conversation) print("Conversation Generation Complete.\n") print("Full Conversation:\n", conversation) # Now we create a separate prompt for the summary summary_prompt = create_summary_prompt(name1, name2, conversation) # Create a new chain for the summary using the summary prompt summary_chain = LLMChain(llm=llm, prompt=summary_prompt) st.subheader("Summary and Title:") st.write("**Summarizing the conversation...**") print("Summarizing the conversation...") # Here we explicitly call the summary chain with the summary prompt # This ensures we are actually sending the summary prompt to the LLM 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()