Update app.py

app.py

@@ -23,13 +23,14 @@ from huggingface_hub import list_datasets, HfApi, hf_hub_download
 ANNOTATION_CONFIG_FILE = "annotation_config.json"
 OUTPUT_FILE_PATH = "dataset.jsonl"
 
-
+llm_provider_state = State("")
 
 def load_llm_config():
     params = load_params()
     return (
         params.get('PROVIDER', ''),
         params.get('BASE_URL', ''),
+        params.get('MODEL', ''),  # Add this line
         params.get('WORKSPACE', ''),
         params.get('API_KEY', ''),
         params.get('max_tokens', 2048),
@@ -41,10 +42,11 @@ def load_llm_config():
 
 
 
-def save_llm_config(provider, base_url, workspace, api_key, max_tokens, temperature, top_p, frequency_penalty, presence_penalty):
+def save_llm_config(provider, base_url, model, workspace, api_key, max_tokens, temperature, top_p, frequency_penalty, presence_penalty):
     save_params({
         'PROVIDER': provider,
         'BASE_URL': base_url,
+        'MODEL': model,  # Add this line
         'WORKSPACE': workspace,
         'API_KEY': api_key,
         'max_tokens': max_tokens,
@@ -56,6 +58,8 @@ def save_llm_config(provider, base_url, workspace, api_key, max_tokens, temperat
     return "LLM configuration saved successfully"
 
 
+def update_model_visibility(provider):
+        return gr.update(visible=provider in ["local-model", "openai"])
 
 
 def load_annotation_config():
@@ -493,7 +497,7 @@ def update_chat_context(row_data, index, total, quality, high_quality_tags, low_
 
 
 
-async def run_generate_dataset(num_workers, num_generations, output_file_path, loaded_dataset):
+async def run_generate_dataset(num_workers, num_generations, output_file_path, llm_provider, dataset):
     if loaded_dataset is None:
         return "Error: No dataset loaded. Please load a dataset before generating.", ""
 
@@ -501,7 +505,7 @@ async def run_generate_dataset(num_workers, num_generations, output_file_path, l
     for _ in range(num_generations):
         topic_selected = random.choice(TOPICS)
         system_message_selected = random.choice(SYSTEM_MESSAGES_VODALUS)
-        data = await generate_data(topic_selected, PROMPT_1, system_message_selected, output_file_path, loaded_dataset)
+        data = await generate_data(topic_selected, PROMPT_1, system_message_selected, output_file_path, llm_provider)
         if data:
             generated_data.append(json.dumps(data))
     
@@ -621,6 +625,13 @@ def load_dataset_wrapper(dataset_name, split):
     dataset, message = load_huggingface_dataset(dataset_name, split)
     return dataset, message
 
+def update_field_visibility(provider):
+    if provider == "local-model":
+        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
+    elif provider == "anything-llm":
+        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
+    else:
+        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
 def get_popular_datasets():
     return [
@@ -913,10 +924,12 @@ with demo:
             with gr.Tab("LLM Configuration"):
                 with gr.Row():
                     provider = gr.Dropdown(choices=["local-model", "anything-llm"], label="LLM Provider")
-                    base_url = gr.Textbox(label="Base URL (for local model)")
                 with gr.Row():
-                    workspace = gr.Textbox(label="Workspace (for AnythingLLM)")
-                    api_key = gr.Textbox(label="API Key (for AnythingLLM)")
+                    base_url = gr.Textbox(label="Base URL (for local model)", visible=False)
+                    model = gr.Textbox(label="Model (for local model)", visible=False)
+                with gr.Row():
+                    workspace = gr.Textbox(label="Workspace (for AnythingLLM)", visible=False)
+                    api_key = gr.Textbox(label="API Key (for AnythingLLM)", visible=False)
                 
                 with gr.Accordion("Advanced Options", open=False):
                     with gr.Row():
@@ -1045,18 +1058,18 @@ with demo:
 
     start_generation_btn.click(
         run_generate_dataset,
-        inputs=[num_workers, num_generations, output_file_path, loaded_dataset],
+        inputs=[num_workers, num_generations, output_file_path, llm_provider, dataset],
         outputs=[generation_status, generation_output]
     )
 
     demo.load(
         load_llm_config,
-        outputs=[provider, base_url, workspace, api_key, max_tokens, temperature, top_p, frequency_penalty, presence_penalty]
+        outputs=[provider, base_url, model, workspace, api_key, max_tokens, temperature, top_p, frequency_penalty, presence_penalty]
     )
 
     save_llm_config_btn.click(
         save_llm_config,
-        inputs=[provider, base_url, workspace, api_key, max_tokens, temperature, top_p, frequency_penalty, presence_penalty],
+        inputs=[provider, base_url, model, workspace, api_key, max_tokens, temperature, top_p, frequency_penalty, presence_penalty],
         outputs=[llm_config_status]
     )
 
@@ -1071,28 +1084,34 @@ with demo:
             outputs=[chatbot]
         )
 
-    search_button.click(
-        search_huggingface_datasets,
-        inputs=[dataset_search],
-        outputs=[dataset_results, dataset_input]
+    provider.change(
+        lambda x: x,
+        inputs=[provider],
+        outputs=[llm_provider_state]
     )
 
-    dataset_results.change(
-        lambda choice: choice,
-        inputs=[dataset_results],
-        outputs=[dataset_input]
-    )
+    # search_button.click(
+    #     search_huggingface_datasets,
+    #     inputs=[dataset_search],
+    #     outputs=[dataset_results, dataset_input]
+    # )
 
-    load_dataset_button.click(
-        load_dataset_wrapper,
-        inputs=[dataset_input, dataset_split],
-        outputs=[loaded_dataset, dataset_status]
-    )
+    # dataset_results.change(
+    #     lambda choice: choice,
+    #     inputs=[dataset_results],
+    #     outputs=[dataset_input]
+    # )
+
+    # load_dataset_button.click(
+    #     load_dataset_wrapper,
+    #     inputs=[dataset_input, dataset_split],
+    #     outputs=[loaded_dataset, dataset_status]
+    # )
 
     # Modify the start_generation_btn.click to include the loaded dataset
     start_generation_btn.click(
         run_generate_dataset,
-        inputs=[num_workers, num_generations, output_file_path, loaded_dataset],
+        inputs=[num_workers, num_generations, output_file_path, llm_provider_state],
         outputs=[generation_status, generation_output]
     )