Update app.py

app.py

@@ -99,7 +99,7 @@ def predict_comprehensive(
     text2: str = "",
     top_k: int = 0,
     include_embeddings: bool = False,
-    include_pca: bool = False,
+    include_layout: bool = False,
     include_unconditional: bool = False,
     use_logprobs: bool = True
 ):
@@ -176,15 +176,13 @@ def predict_comprehensive(
             "first_10_tokens_sample": embeddings_cache[:10, :10].tolist() if embeddings_cache is not None else None
         }
     
-    # PCA projections if requested
-    if include_pca:
+    # UMAP layout if requested
+    if include_layout:
         if umap_projections_cache is not None:
-            # For compatibility we keep the `pca` key but fill with UMAP projections
-            result["pca"] = {
-                "projections": umap_projections_cache.tolist(),
-                "explained_variance_ratio": [0.0, 0.0],  # Placeholder; UMAP has no variance ratio
+            result["layout"] = {
                 "method": "umap",
                 "umap_params": umap_params,
+                "projections": umap_projections_cache.tolist(),
             }
     
     return result
@@ -219,12 +217,12 @@ with gr.Blocks(title="Token Probability Visualization API") as demo:
     gr.Markdown("""
     # Token Probability Visualization API
     
-    This API provides token probabilities, embeddings, and PCA projections for visualization.
+    This API provides token probabilities, embeddings, and a precomputed UMAP layout for visualization.
     Model: `openai-community/gpt2`
     """)
     
     with gr.Tab("Comprehensive API"):
-        gr.Markdown("### Get token probabilities with optional embeddings and PCA")
+        gr.Markdown("### Get token probabilities with optional embeddings and 2D layout")
         
         with gr.Row():
             with gr.Column():
@@ -242,7 +240,7 @@ with gr.Blocks(title="Token Probability Visualization API") as demo:
                 with gr.Row():
                     top_k_slider = gr.Slider(0, 200, step=1, value=20, label="Top-K tokens")
                     include_embeddings = gr.Checkbox(False, label="Include Embeddings Sample")
-                    include_pca = gr.Checkbox(False, label="Include UMAP Projections (2D)")
+                    include_layout = gr.Checkbox(False, label="Include 2D Layout (UMAP)")
                     include_unconditional = gr.Checkbox(True, label="Include Unconditional")
                     use_logprobs = gr.Checkbox(True, label="Include Log Probabilities")
                 
@@ -253,8 +251,8 @@ with gr.Blocks(title="Token Probability Visualization API") as demo:
         
         predict_btn.click(
             fn=predict_comprehensive,
-            inputs=[text_input, text2_input, top_k_slider, include_embeddings, 
-                   include_pca, include_unconditional, use_logprobs],
+            inputs=[text_input, text2_input, top_k_slider, include_embeddings,
+                   include_layout, include_unconditional, use_logprobs],
             outputs=output_json,
             api_name="predict"
         )