Update app.py

app.py

@@ -10,6 +10,7 @@ from main import generate_data, PROMPT_1
 from topics import TOPICS
 from system_messages import SYSTEM_MESSAGES_VODALUS
 import random
+import spaces
 
 ANNOTATION_CONFIG_FILE = "annotation_config.json"
 OUTPUT_FILE_PATH = "dataset.jsonl"
@@ -263,7 +264,7 @@ def save_dataset_config(system_messages, prompt_1, topics):
     
     return "Dataset configuration saved successfully"
 
-# Modify the chat_with_llm function to use Gradio's built-in async capabilities
+@spaces.GPU(duration=120)
 def chat_with_llm(message, history):
     try:
         msg_list = [{"role": "system", "content": "You are an AI assistant helping with dataset annotation and quality checking."}]
@@ -291,7 +292,7 @@ def update_chat_context(row_data, index, total, quality, high_quality_tags, low_
     """
     return [[None, context]]  # Return as a list of message pairs
 
-# Add this function to handle dataset generation
+@spaces.GPU(duration=180)
 async def run_generate_dataset(num_workers, num_generations, output_file_path):
     generated_data = []
     for _ in range(num_generations):