Update app.py

app.py

@@ -7,7 +7,6 @@ import numpy as np
 import streamlit as st
 from openai import OpenAI
 import os
-import sys
 from dotenv import load_dotenv
 
 load_dotenv()
@@ -23,25 +22,14 @@ 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
 
-# Define the available models
-models = [key for key in model_links.keys()]
-
 # Sidebar for model selection
-selected_model = st.sidebar.selectbox("Select Model", models)
+selected_model = st.sidebar.selectbox("Select Model", list(model_links.keys()))
 
 # Temperature slider
 temp_values = st.sidebar.slider('Select a temperature value', 0.0, 1.0, 0.5)
@@ -115,29 +103,48 @@ if task_choice == "Data Generation":
     st.code(system_prompt)
 
     if st.button("Generate Examples"):
-        # Generate examples by concatenating all inputs and sending it to the model
+        # Break generation into smaller chunks if needed to ensure enough examples
+        all_generated_examples = []
+        remaining_examples = num_to_generate
+
         with st.spinner("Generating..."):
-            st.session_state.messages.append({"role": "system", "content": system_prompt})
-
-            try:
-                stream = client.chat.completions.create(
-                    model=model_links[selected_model],
-                    messages=[
-                        {"role": m["role"], "content": m["content"]}
-                        for m in st.session_state.messages
-                    ],
-                    temperature=temp_values,
-                    stream=True,
-                    max_tokens=3000,
-                )
-                response = st.write_stream(stream)
-            except Exception as e:
-                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})
+            while remaining_examples > 0:
+                # Generating examples in chunks to avoid token limits
+                chunk_size = min(remaining_examples, 5)
+                try:
+                    # Append the new chunk system prompt
+                    st.session_state.messages.append({"role": "system", "content": system_prompt})
+
+                    stream = client.chat.completions.create(
+                        model=model_links[selected_model],
+                        messages=[
+                            {"role": m["role"], "content": m["content"]}
+                            for m in st.session_state.messages
+                        ],
+                        temperature=temp_values,
+                        stream=True,
+                        max_tokens=3000,
+                    )
+
+                    # Parse the response and break it into individual examples
+                    response = st.write_stream(stream)
+                    generated_examples = response.split("\n")[:chunk_size]  # Assuming each example is on a new line
+
+                    # Store the new examples
+                    all_generated_examples.extend(generated_examples)
+                    remaining_examples -= chunk_size
+
+                except Exception as e:
+                    st.error("Error during generation.")
+                    st.write(e)
+                    break
+
+        # Display all generated examples
+        for idx, example in enumerate(all_generated_examples):
+            st.write(f"Example {idx+1}: {example}")
+
+        # Update session state to prevent repetition of old prompts
+        st.session_state.messages = []  # Clear messages after each generation
 
 else:
     # Data labeling workflow (for future implementation based on classification)
@@ -145,4 +152,3 @@ else:
 
 
 
-