Update app.py

app.py

@@ -10,6 +10,7 @@ import numpy as np
 import streamlit as st
 from openai import OpenAI
 import os
+import sys
 from dotenv import load_dotenv
 
 load_dotenv()
@@ -17,124 +18,131 @@ load_dotenv()
 # Initialize the client
 client = OpenAI(
     base_url="https://api-inference.huggingface.co/v1",
-    #api_key=os.environ.get('HUGGINGFACEHUB_API_TOKEN')  # Replace with your token
-    api_key=os.environ.get('HF_TOKEN')  # Replace with your token
-
+    api_key=os.environ.get('HUGGINGFACEHUB_API_TOKEN')  # Add your Huggingface token here
 )
 
-# Function to reset conversation
+# Supported models
+model_links = {
+    "Meta-Llama-3-8B": "meta-llama/Meta-Llama-3-8B-Instruct"
+}
+
+# Random dog images for error messages
+random_dog = [
+    "0f476473-2d8b-415e-b944-483768418a95.jpg",
+    "1bd75c81-f1d7-4e55-9310-a27595fa8762.jpg",
+    "526590d2-8817-4ff0-8c62-fdcba5306d02.jpg",
+    "1326984c-39b0-492c-a773-f120d747a7e2.jpg"
+]
+
+# Reset conversation
 def reset_conversation():
     st.session_state.conversation = []
     st.session_state.messages = []
     return None
 
-# Initialize session state for 'messages' if it doesn't exist
-if 'messages' not in st.session_state:
-    st.session_state.messages = []
+# Define the available models
+models = [key for key in model_links.keys()]
 
-# Define classification options
-classification_types = ["Sentiment Analysis", "Binary Classification", "Multi-Class Classification"]
+# Sidebar for model selection
+selected_model = st.sidebar.selectbox("Select Model", models)
 
-# Start with a selection between data generation or labeling
-st.sidebar.write("Choose Task:")
-task = st.sidebar.radio("Do you want to generate data or label data?", ("Data Generation", "Data Labeling"))
+# Temperature slider
+temp_values = st.sidebar.slider('Select a temperature value', 0.0, 1.0, 0.5)
 
-# If the user selects Data Labeling
-if task == "Data Labeling":
-    st.sidebar.write("Choose Classification Type:")
-    classification_type = st.sidebar.radio("Select a classification type:", classification_types)
+# Reset button
+st.sidebar.button('Reset Chat', on_click=reset_conversation)
 
-    # Handle Sentiment Analysis
-    if classification_type == "Sentiment Analysis":
-        st.sidebar.write("Classes: Positive, Negative, Neutral (fixed)")
-        class_labels = ["Positive", "Negative", "Neutral"]
+# Model description
+st.sidebar.write(f"You're now chatting with **{selected_model}**")
+st.sidebar.markdown("*Generated content may be inaccurate or false.*")
 
-    # Handle Binary Classification
-    elif classification_type == "Binary Classification":
-        class_1 = st.sidebar.text_input("Enter Class 1:")
-        class_2 = st.sidebar.text_input("Enter Class 2:")
-        class_labels = [class_1, class_2]
-
-    # Handle Multi-Class Classification
-    elif classification_type == "Multi-Class Classification":
-        class_labels = []
-        for i in range(1, 11):  # Allow up to 10 classes
-            label = st.sidebar.text_input(f"Enter Class {i} (leave blank to stop):")
-            if label:
-                class_labels.append(label)
-            else:
-                break
-
-    # Domain selection
-    st.sidebar.write("Specify the Domain:")
-    domain = st.sidebar.radio("Choose a domain:", ("Restaurant Reviews", "E-commerce Reviews", "Custom"))
-    if domain == "Custom":
-        domain = st.sidebar.text_input("Enter Custom Domain:")
-
-    # Specify example length
-    st.sidebar.write("Specify the Length of Examples:")
-    min_words = st.sidebar.number_input("Minimum word count (10 to 90):", 10, 90, 10)
-    max_words = st.sidebar.number_input("Maximum word count (10 to 90):", min_words, 90, 50)
-
-    # Few-shot examples option
-    use_few_shot = st.sidebar.radio("Do you want to use few-shot examples?", ("Yes", "No"))
-    few_shot_examples = []
-    if use_few_shot == "Yes":
-        num_examples = st.sidebar.number_input("How many few-shot examples? (1 to 5)", 1, 5, 1)
-        for i in range(num_examples):
-            example_text = st.text_area(f"Enter example {i+1}:")
-            example_label = st.selectbox(f"Select the label for example {i+1}:", class_labels)
-            few_shot_examples.append({"text": example_text, "label": example_label})
-
-    # Generate the system prompt based on classification type
-    if classification_type == "Sentiment Analysis":
-        system_prompt = f"You are a propositional sentiment analysis expert. Your role is to generate sentiment analysis reviews based on the data entered and few-shot examples provided, if any, for the domain '{domain}'."
-    elif classification_type == "Binary Classification":
-        system_prompt = f"You are an expert in binary classification. Your task is to label examples for the domain '{domain}' with either '{class_1}' or '{class_2}', based on the data provided."
-    else:  # Multi-Class Classification
-        system_prompt = f"You are an expert in multi-class classification. Your role is to label examples for the domain '{domain}' using the provided class labels."
+# Chat initialization
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
-    st.sidebar.write("System Prompt:")
-    st.sidebar.write(system_prompt)
+# Display chat messages
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
 
-    # Step-by-step thinking
-    st.sidebar.write("Generated Data:")
-    st.sidebar.write("Think step by step to ensure accuracy in classification.")
+# Main logic to choose between data generation and data labeling
+task_choice = st.selectbox("Choose Task", ["Data Generation", "Data Labeling"])
 
-    # Accept user input for generating or labeling data
-    if prompt := st.chat_input(f"Hi, I'm ready to help with {classification_type} for {domain}. Ask me a question or provide data to classify."):
+if task_choice == "Data Generation":
+    classification_type = st.selectbox(
+        "Choose Classification Type",
+        ["Sentiment Analysis", "Binary Classification", "Multi-Class Classification"]
+    )
 
-        # Display user message in chat message container
-        with st.chat_message("user"):
-            st.markdown(prompt)
-        # Add user message to chat history
-        st.session_state.messages.append({"role": "user", "content": prompt})
+    if classification_type == "Sentiment Analysis":
+        st.write("Sentiment Analysis: Positive, Negative, Neutral")
+        labels = ["Positive", "Negative", "Neutral"]
+    elif classification_type == "Binary Classification":
+        label_1 = st.text_input("Enter first class")
+        label_2 = st.text_input("Enter second class")
+        labels = [label_1, label_2]
+    elif classification_type == "Multi-Class Classification":
+        num_classes = st.slider("How many classes?", 3, 10, 3)
+        labels = [st.text_input(f"Class {i+1}") for i in range(num_classes)]
 
-        # Display assistant response in chat message container
-        with st.chat_message("assistant"):
+    domain = st.selectbox("Choose Domain", ["Restaurant reviews", "E-commerce reviews", "Custom"])
+    if domain == "Custom":
+        domain = st.text_input("Specify custom domain")
+
+    min_words = st.number_input("Minimum words per example", min_value=10, max_value=90, value=10)
+    max_words = st.number_input("Maximum words per example", min_value=10, max_value=90, value=90)
+
+    few_shot = st.radio("Do you want to use few-shot examples?", ["Yes", "No"])
+    if few_shot == "Yes":
+        num_examples = st.slider("How many few-shot examples?", 1, 5, 1)
+        few_shot_examples = [
+            {"content": st.text_area(f"Example {i+1}"), "label": st.selectbox(f"Label for example {i+1}", labels)}
+            for i in range(num_examples)
+        ]
+    else:
+        few_shot_examples = []
+
+    # Ask the user how many examples they need
+    num_to_generate = st.number_input("How many examples to generate?", min_value=1, max_value=50, value=10)
+
+    # System prompt generation
+    system_prompt = f"You are a professional {classification_type.lower()} expert. Your role is to generate data for {domain}.\n\n"
+    if few_shot_examples:
+        system_prompt += "Use the following few-shot examples as a reference:\n"
+        for example in few_shot_examples:
+            system_prompt += f"Example: {example['content']}, Label: {example['label']}\n"
+    system_prompt += f"Each example should have between {min_words} and {max_words} words.\n"
+    system_prompt += "Think step by step while generating the examples."
+
+    st.write("System Prompt:")
+    st.code(system_prompt)
+
+    if st.button("Generate Examples"):
+        # Generate examples by concatenating all inputs and sending it to the model
+        with st.spinner("Generating..."):
+            st.session_state.messages.append({"role": "system", "content": system_prompt})
 
             try:
-                # Stream the response from the model
                 stream = client.chat.completions.create(
-                    model="meta-llama/Meta-Llama-3-8B-Instruct",
+                    model=model_links[selected_model],
                     messages=[
                         {"role": m["role"], "content": m["content"]}
                         for m in st.session_state.messages
                     ],
-                    temperature=0.5,
+                    temperature=temp_values,
                     stream=True,
                     max_tokens=3000,
                 )
-
                 response = st.write_stream(stream)
-
             except Exception as e:
-                response = "😵‍💫 Something went wrong. Try again later."
-                st.write(response)
+                response = "Error during generation."
+                random_dog_pick = 'https://random.dog/' + random_dog[np.random.randint(len(random_dog))]
+                st.image(random_dog_pick)
+                st.write(e)
 
-        st.session_state.messages.append({"role": "assistant", "content": response})
+            st.session_state.messages.append({"role": "assistant", "content": response})
 
-# If the user selects Data Generation
 else:
-    st.sidebar.write("This feature will allow you to generate new data. Coming soon!")
+    # Data labeling workflow (for future implementation based on classification)
+    st.write("Data Labeling functionality will go here.")