Update app.py

app.py

@@ -26,22 +26,41 @@ class StringStoppingCriteria(StoppingCriteria):
 def create_demo(api_url, model_id=None):
     # Standalone mode check
     hf_model_id = model_id or os.environ.get("HF_MODEL_ID", "ethicalabs/Echo-DSRN-Small-Instruct")
-    
+
     local_model = None
     local_tokenizer = None
-    
+    model_metadata = None
+
     if hf_model_id:
         print(f"📦 Standalone Mode: Loading {hf_model_id} via Transformers...")
         try:
             local_tokenizer = AutoTokenizer.from_pretrained(hf_model_id, trust_remote_code=True)
             local_model = AutoModelForCausalLM.from_pretrained(
-                hf_model_id, 
+                hf_model_id,
                 trust_remote_code=True,
                 torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-                device_map="auto"
+                device_map="auto",
             )
             local_model.eval()
             print("✅ Model loaded successfully.")
+
+            # Calculate metadata
+            total_params = sum(p.numel() for p in local_model.parameters())
+            trainable_params = sum(p.numel() for p in local_model.parameters() if p.requires_grad)
+            config = local_model.config
+
+            model_metadata = {
+                "Total Parameters": f"{total_params:,}",
+                "Trainable Parameters": f"{trainable_params:,}",
+                "Layers": getattr(
+                    config, "num_hidden_layers", getattr(config, "num_layers", "N/A")
+                ),
+                "Hidden Size": getattr(config, "hidden_size", "N/A"),
+                "MLP Ratio": getattr(config, "mlp_ratio", "N/A"),
+                "Vocab Size": getattr(config, "vocab_size", "N/A"),
+                "Surprise Lambda Init": getattr(config, "surprise_lambda_init", "N/A"),
+                "Model Type": getattr(config, "model_type", "echo"),
+            }
         except Exception as e:
             print(f"❌ Failed to load model: {e}")
             traceback.print_exc()
@@ -51,43 +70,49 @@ def create_demo(api_url, model_id=None):
         font-family: 'JetBrains Mono', 'Courier New', monospace;
         line-height: 1.8;
         padding: 24px;
-        background: #0f172a;
+        background: #020617;
         border-radius: 12px;
         color: #e2e8f0;
         min-height: 250px;
         white-space: pre-wrap;
         overflow-wrap: break-word;
         border: 1px solid #1e293b;
+        position: relative;
+        overflow: visible !important;
     }
     .token {
-        display: inline;
-        padding: 1px 2px;
-        border-radius: 4px;
+        display: inline-block;
+        padding: 0 1px;
+        border-radius: 3px;
         position: relative;
         cursor: help;
-        transition: background 0.2s;
+        transition: all 0.2s;
     }
     .token:hover {
-        background: rgba(255, 255, 255, 0.1);
+        background: rgba(255, 255, 255, 0.15) !important;
+        z-index: 100;
     }
     .token-tooltip {
-        visibility: hidden;
+        display: none;
         position: absolute;
-        bottom: 120%;
+        bottom: 130%;
         left: 50%;
         transform: translateX(-50%);
         background: #1e293b;
         color: white;
-        padding: 6px 10px;
-        border-radius: 6px;
-        font-size: 0.8rem;
-        z-index: 100;
+        padding: 12px;
+        border-radius: 10px;
+        font-size: 0.85rem;
+        z-index: 1000;
         white-space: nowrap;
-        box-shadow: 0 4px 12px rgba(0,0,0,0.5);
-        border: 1px solid #334155;
+        box-shadow: 0 20px 25px -5px rgba(0, 0, 0, 0.5);
+        border: 1px solid rgba(255, 255, 255, 0.2);
+        pointer-events: none;
+        text-align: center;
+        min-width: 120px;
     }
     .token:hover .token-tooltip {
-        visibility: visible;
+        display: block;
     }
     .mode-badge {
         display: inline-block;
@@ -99,21 +124,46 @@ def create_demo(api_url, model_id=None):
     }
     .mode-standalone { background: #059669; color: white; }
     .mode-api { background: #4f46e5; color: white; }
+    
+    .metadata-card {
+        background: #0f172a;
+        padding: 15px;
+        border-radius: 8px;
+        border: 1px solid #1e293b;
+        font-family: inherit;
+    }
+    .metadata-grid {
+        display: grid;
+        grid-template-columns: repeat(auto-fill, minmax(180px, 1fr));
+        gap: 10px;
+        margin-top: 10px;
+    }
+    .metadata-item {
+        font-size: 0.85rem;
+        color: #94a3b8;
+    }
+    .metadata-value {
+        font-weight: 600;
+        color: #f1f5f9;
+        display: block;
+    }
     """
 
-    def generate_html(prompt, max_tokens, temperature, top_k):
+    def generate_html(prompt, max_tokens, temperature, top_k, rep_penalty):
         try:
             if local_model and local_tokenizer:
                 # Standalone Mode Logic
                 input_tokens = local_tokenizer(prompt, return_tensors="pt").to(local_model.device)
                 input_ids = input_tokens.input_ids
-                
+
                 stop_strings = ["<|im_end|>", "<|end|>", "<|user|>"]
-                stopping_criteria = StoppingCriteriaList([StringStoppingCriteria(local_tokenizer, stop_strings)])
-                
+                stopping_criteria = StoppingCriteriaList(
+                    [StringStoppingCriteria(local_tokenizer, stop_strings)]
+                )
+
                 all_tokens = []
                 all_logprobs = []
-                
+
                 # Echo (Prompt) Logprobs
                 with torch.no_grad():
                     outputs = local_model(input_ids)
@@ -122,13 +172,13 @@ def create_demo(api_url, model_id=None):
                         token_id = input_ids[0, i].item()
                         token_text = local_tokenizer.decode([token_id])
                         all_tokens.append(token_text)
-                        
+
                         if i == 0:
                             all_logprobs.append(None)
                         else:
-                            lp = torch.nn.functional.log_softmax(logits[0, i-1, :], dim=-1)
+                            lp = torch.nn.functional.log_softmax(logits[0, i - 1, :], dim=-1)
                             all_logprobs.append(lp[token_id].item())
-                
+
                 # Generation
                 with torch.no_grad():
                     gen_out = local_model.generate(
@@ -137,27 +187,27 @@ def create_demo(api_url, model_id=None):
                         temperature=temperature,
                         top_k=top_k,
                         top_p=0.9,
-                        repetition_penalty=1.2,
+                        repetition_penalty=rep_penalty,
                         do_sample=temperature > 0,
                         use_cache=False,
                         output_scores=True,
                         return_dict_in_generate=True,
                         stopping_criteria=stopping_criteria,
                         pad_token_id=local_tokenizer.pad_token_id or 32000,
-                        eos_token_id=[32000, 32007, 32011]
+                        eos_token_id=[32000, 32007, 32011],
                     )
-                
+
                 output_ids = gen_out.sequences
                 scores = gen_out.scores
-                generated_ids = output_ids[0, input_ids.shape[1]:]
-                
+                generated_ids = output_ids[0, input_ids.shape[1] :]
+
                 for i, token_id in enumerate(generated_ids):
                     token_id = token_id.item()
                     token_text = local_tokenizer.decode([token_id])
                     all_tokens.append(token_text)
                     lp = torch.nn.functional.log_softmax(scores[i][0, :], dim=-1)
                     all_logprobs.append(lp[token_id].item())
-                
+
                 logprobs = {"tokens": all_tokens, "token_logprobs": all_logprobs}
             else:
                 # API Mode
@@ -166,6 +216,7 @@ def create_demo(api_url, model_id=None):
                     "max_tokens": max_tokens,
                     "temperature": temperature,
                     "top_k": top_k,
+                    "repetition_penalty": rep_penalty,
                     "logprobs": 1,
                     "echo": True,
                 }
@@ -190,14 +241,21 @@ def create_demo(api_url, model_id=None):
                 color = "#e2e8f0"  # Default
                 if lp is not None:
                     if prob > 0.8:
-                        color = "#4ade80"  # Green
+                        color = "#4ade80"  # Vibrant Green
                     elif prob < 0.2:
-                        color = "#f87171"  # Red
+                        color = "#f87171"  # Soft Red
 
                 prob_pct = f"{(prob * 100):.1f}%" if lp is not None else "N/A"
+                logit_str = f"{lp:.4f}" if lp is not None else "N/A"
 
                 html += f'<span class="token" style="color: {color}">{display_token}'
-                html += f'<span class="token-tooltip">Prob: {prob_pct}</span></span>'
+                html += f'<span class="token-tooltip">'
+                html += f'<span style="opacity: 0.7; font-size: 0.7rem">CONFIDENCE</span><br>'
+                html += f'<span style="font-size: 1.1rem; font-weight: 800; color: {color}">{prob_pct}</span><br>'
+                html += f'<hr style="margin: 5px 0; opacity: 0.1; border: none; border-top: 1px solid white;">'
+                html += f'<span style="opacity: 0.7; font-size: 0.7rem">LOGIT</span><br>'
+                html += f'<span style="font-family: monospace">{logit_str}</span>'
+                html += f'</span></span>'
 
             html += "</div>"
             return html
@@ -205,22 +263,38 @@ def create_demo(api_url, model_id=None):
             traceback.print_exc()
             return f'<div class="token-container" style="color: #f87171">❌ Error: {str(e)}</div>'
 
-    with gr.Blocks(title="Echo-DSRN Predictor") as demo:
-        gr.Markdown(
-            """
-            # 🌫️ Echo-DSRN Next Word Prediction
-            ### Responsible AI Development by ethicalabs.ai
-            """
-        )
+    with gr.Blocks(title="Echo-DSRN-Small-400m - Next Word Prediction") as demo:
+        with gr.Row():
+            with gr.Column(scale=4):
+                gr.Markdown(
+                    """
+                    # 🌫️ Echo-DSRN-Small-400m - Next Word Prediction
+                    ### Responsible AI Development by ethicalabs.ai
+                    
+                    **Echo-DSRN** is a high-performance **Dual-State Recurrent Network** inspired by the architecture in [Titans: Learning to Memorize at Test Time (Google Research)](https://arxiv.org/abs/2501.00663). It combines the efficiency of linear recurrence with optimized gating and a novel **Surprise Mechanism** to enable infinite context extrapolation.
+                    
+                    <p style="margin-top: 15px; border-top: 1px solid rgba(255,255,255,0.1); padding-top: 10px;">
+                    <span style="font-weight: bold; color: #f87171;">⚠️ Proprietary Technology Disclaimer</span><br>
+                    <span style="font-style: italic; font-size: 0.9rem; color: #94a3b8;">This research is currently in an internal-only alpha phase. The model source code and weights can be released only when we have the capacity to scale up and expand and fine-tune and safety align it according to our strict responsibility standards.</span>
+                    </p>
+                    """
+                )
+            with gr.Column(scale=1):
+                mode_str = "Standalone Mode" if local_model else f"Connected to API: {api_url}"
+                mode_cls = "mode-standalone" if local_model else "mode-api"
+                gr.HTML(f'<div class="mode-badge {mode_cls}">{mode_str}</div>')
 
-        mode_str = "Standalone Mode" if local_model else f"Connected to API: {api_url}"
-        mode_cls = "mode-standalone" if local_model else "mode-api"
-        gr.HTML(f'<div class="mode-badge {mode_cls}">{mode_str}</div>')
+        if model_metadata:
+            gr.Markdown("### 📦 Model Architecture Details")
+            metadata_html = '<div class="metadata-card"><div class="metadata-grid">'
+            for key, value in model_metadata.items():
+                metadata_html += f'<div class="metadata-item">{key}<span class="metadata-value">{value}</span></div>'
+            metadata_html += "</div></div>"
+            gr.HTML(metadata_html)
 
         gr.Markdown(
             """
-            This application shows the internal probabilities of the model. 
-            Hover over tokens to see their prediction confidence.
+            Hover over generated tokens to visualize the model's confidence and raw logit values.
             """
         )
 
@@ -238,6 +312,9 @@ def create_demo(api_url, model_id=None):
                     )
                 with gr.Row():
                     top_k = gr.Slider(minimum=0, maximum=100, value=40, step=1, label="Top K")
+                    rep_penalty = gr.Slider(
+                        minimum=1.0, maximum=2.0, value=1.2, step=0.05, label="Repetition Penalty"
+                    )
 
                 submit_btn = gr.Button("Generate with Probabilities", variant="primary")
 
@@ -246,16 +323,18 @@ def create_demo(api_url, model_id=None):
                 output_html = gr.HTML(label="Visualized Output")
 
         submit_btn.click(
-            generate_html, inputs=[prompt, max_tokens, temp, top_k], outputs=output_html
+            generate_html,
+            inputs=[prompt, max_tokens, temp, top_k, rep_penalty],
+            outputs=output_html,
         )
 
         gr.Examples(
             examples=[
-                ["The capital of France is", 128, 0.1, 40],
-                ["Responsible AI development requires", 128, 0.7, 40],
-                ["To be or not to be, that is the", 128, 0.1, 40],
+                ["The capital of France is", 128, 0.1, 40, 1.2],
+                ["Responsible AI development requires", 128, 0.7, 40, 1.2],
+                ["To be or not to be, that is the", 128, 0.1, 40, 1.2],
             ],
-            inputs=[prompt, max_tokens, temp, top_k],
+            inputs=[prompt, max_tokens, temp, top_k, rep_penalty],
         )
 
     return demo, css
@@ -267,10 +346,10 @@ if __name__ == "__main__":
         "--api_url", type=str, default="http://localhost:5000", help="URL of the Echo-DSRN API"
     )
     parser.add_argument(
-        "--model_id", 
-        type=str, 
-        default="ethicalabs/Echo-DSRN-Small-Instruct", 
-        help="Local path or HF ID to the model for Standalone mode"
+        "--model_id",
+        type=str,
+        default=None,
+        help="Local path or HF ID to the model for Standalone mode",
     )
     parser.add_argument("--port", type=int, default=7860, help="Port to run the Gradio app on")
     args = parser.parse_args()