Update app.py

app.py

@@ -233,13 +233,13 @@ def display_confusion_matrices_los(percentage):
     return raw_img, embeddings_img
 
 # Main function to handle user choice
-def handle_user_choice(choice, percentage=None, uploaded_file=None):
+def handle_user_choice(choice, percentage=None, uploaded_file=None, emb_type='CLS Embedding'):
     if choice == "Use Default Dataset":
         raw_img, embeddings_img = display_confusion_matrices_los(percentage)
         return raw_img, embeddings_img, ""  # Return empty string for console output
     elif choice == "Upload Dataset":
         if uploaded_file is not None:
-            raw_img, embeddings_img, console_output = process_hdf5_file(uploaded_file, percentage)
+            raw_img, embeddings_img, console_output = process_hdf5_file(uploaded_file, percentage, emb_type)
             return raw_img, embeddings_img, console_output
         else:
             return "Please upload a dataset", "Please upload a dataset", ""  # Return empty string for console output
@@ -409,7 +409,7 @@ def identical_train_test_split(output_emb, output_raw, labels, train_percentage)
 # Store the original working directory when the app starts
 original_dir = os.getcwd()
 
-def process_hdf5_file(uploaded_file, percentage):
+def process_hdf5_file(uploaded_file, percentage, emb_type='CLS Embedding'):
     capture = PrintCapture()
     sys.stdout = capture  # Redirect print statements to capture
     
@@ -462,7 +462,12 @@ def process_hdf5_file(uploaded_file, percentage):
         preprocessed_chs = input_preprocess.tokenizer(manual_data=channels)
 
         # Step 7: Perform inference using the functions from inference.py
-        output_emb = inference.lwm_inference(preprocessed_chs, 'channel_emb', model, device)
+        if emb_type == 'Channel Embedding':
+            embedding_type = 'channel_emb'
+        elif emb_type == 'CLS Embedding':
+            embedding_type = 'cls_emb'
+            
+        output_emb = inference.lwm_inference(preprocessed_chs, embedding_type, model, device)
         output_raw = inference.create_raw_dataset(preprocessed_chs, device)
 
         print(f"Output Embeddings Shape: {output_emb.shape}")
@@ -648,6 +653,10 @@ with gr.Blocks(css="""
             embeddings_img_los = gr.Image(label="Embeddings", type="pil", width=300, height=300)
             output_textbox = gr.Textbox(label="Console Output", lines=10, elem_id="console-output")
 
+            emb_type = gr.Dropdown(choices=["Channel Embedding", "CLS Embedding"], 
+                           value="CLS Embedding", 
+                           label="Embedding Type", interactive=True)
+            
         # Update the file uploader visibility based on user choice
         def toggle_file_input(choice):
             return gr.update(visible=(choice == "Upload Dataset"))
@@ -655,13 +664,17 @@ with gr.Blocks(css="""
         choice_radio.change(fn=toggle_file_input, inputs=[choice_radio], outputs=file_input)
 
         # When user makes a choice, update the display
-        choice_radio.change(fn=handle_user_choice, inputs=[choice_radio, percentage_slider_los, file_input], 
+        choice_radio.change(fn=handle_user_choice, inputs=[choice_radio, percentage_slider_los, file_input, emb_type], 
                             outputs=[raw_img_los, embeddings_img_los, output_textbox])
 
         # When percentage slider changes (for predefined data)
-        percentage_slider_los.change(fn=handle_user_choice, inputs=[choice_radio, percentage_slider_los, file_input], 
+        percentage_slider_los.change(fn=handle_user_choice, inputs=[choice_radio, percentage_slider_los, file_input, emb_type], 
                                        outputs=[raw_img_los, embeddings_img_los, output_textbox])
 
+        # When embedding type changes
+        emb_type.change(fn=handle_user_choice, inputs=[choice_radio, percentage_slider_los, file_input, emb_type], 
+                        outputs=[raw_img_los, embeddings_img_los, output_textbox])
+
         # Add a conclusion section at the bottom
         gr.Markdown("""
             <div class="explanation-box">