Update app.py

app.py

@@ -1,107 +1,109 @@
 import gradio as gr
 import torch
 import yaml
-import os
+import json
+from tokenizers import Tokenizer
 
-def load_model():
-    """Load the HRM model and config"""
-    try:
-        # Load config
-        with open('all_config.yaml', 'r') as f:
-            config = yaml.safe_load(f)
-        
-        # Load checkpoint
-        checkpoint = torch.load('pytorch_model.bin', map_location='cpu')
-        
-        return config, checkpoint, "✅ Model loaded successfully!"
-    except Exception as e:
-        return None, None, f"❌ Error loading model: {str(e)}"
+# --- 1. Load Custom Model Code ---
+# This dynamically loads your corrected HRM source code
+from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1
 
-def test_model_info(config, checkpoint):
-    """Display model information"""
-    if config is None or checkpoint is None:
-        return "Model not loaded"
-    
-    info = f"""
-    **Model Architecture**: {config['arch']['name']}
-    **Hidden Size**: {config['arch']['hidden_size']}
-    **H Layers**: {config['arch']['H_layers']}
-    **L Layers**: {config['arch']['L_layers']}
-    **Parameters in Checkpoint**: {len(checkpoint)}
-    **Model Purpose**: Grant Abstract Optimization
-    
-    **Training Details**:
-    - Steps: 492,500 (final checkpoint)
-    - Batch Size: {config['global_batch_size']}
-    - Learning Rate: {config['lr']}
-    """
-    return info
+# --- 2. Load Artifacts ---
+print("Loading artifacts...")
+# Load the tokenizer
+tokenizer = Tokenizer.from_file("tokenizer.json")
+# Load the model configuration
+with open('config.yaml', 'r') as f:
+    config_data = yaml.safe_load(f)
+model_config = config_data['arch']
+# Load the grant type mapping
+with open('activity_code_map.json', 'r') as f:
+    activity_code_map = json.load(f)
 
-def placeholder_inference(draft_abstract, grant_type):
-    """Placeholder for actual inference (requires full training pipeline)"""
-    return f"""
-    **Input Abstract**: {draft_abstract[:100]}...
-    
-    **Grant Type**: {grant_type}
-    
-    **Status**: Model checkpoint loaded successfully! 
-    
-    ⚠️ **Note**: Full inference requires the original training pipeline with tokenizer and preprocessing code. 
-    This demo shows that the model weights are accessible and the architecture is properly configured.
-    
-    **Next Steps**: 
-    1. Integrate with original training codebase
-    2. Load tokenizer and preprocessing pipeline  
-    3. Implement full inference function
+# --- 3. Initialize the Model ---
+print("Initializing model...")
+# The model expects a dict, so we pass the Pydantic model's dict representation
+# We also need to add other required keys from the root of the config
+model_config.update({
+    'batch_size': config_data['global_batch_size'],
+    'seq_len': 512, # You may need to get this from your dataset metadata
+    'num_puzzle_identifiers': len(activity_code_map) + 1,
+    'vocab_size': tokenizer.get_vocab_size()
+})
+model = HierarchicalReasoningModel_ACTV1(config_dict=model_config)
+# Load the fine-tuned weights
+model.load_state_dict(torch.load('pytorch_model.bin', map_location='cpu'))
+model.eval() # Set the model to evaluation mode
+print("Model loaded successfully!")
+
+# --- 4. Define the Inference Function ---
+def optimize_abstract(draft_abstract, grant_type):
+    """
+    Takes a draft abstract and grant type, runs the model, and returns the optimized text.
     """
+    if not draft_abstract or not grant_type:
+        return "Please provide both a draft abstract and a grant type."
 
-# Load model on startup
-config, checkpoint, load_status = load_model()
+    try:
+        # Prepare inputs
+        tokenizer.enable_padding(length=512)
+        tokenizer.enable_truncation(max_length=512)
+        
+        input_ids = tokenizer.encode(draft_abstract).ids
+        grant_type_id = activity_code_map.get(grant_type, 0) # Default to 0 if unknown
 
-# Create Gradio interface
-with gr.Blocks(title="HRM Grant Abstract Optimizer") as demo:
-    gr.Markdown("# 🎯 Hierarchical Reasoning Model for Grant Abstract Optimization")
-    gr.Markdown("A specialized 27M-parameter model for transforming draft grant abstracts into funding-worthy versions.")
-    
-    with gr.Tab("Model Info"):
-        gr.Markdown("## Model Status")
-        gr.Markdown(load_status)
+        # Convert to PyTorch tensors
+        input_tensor = torch.tensor([input_ids], dtype=torch.long)
+        grant_tensor = torch.tensor([grant_type_id], dtype=torch.long)
+
+        # Create the batch dictionary that the model expects
+        batch = {
+            "inputs": input_tensor,
+            "puzzle_identifiers": grant_tensor,
+            # The model requires a 'labels' field, even for inference, so we provide a dummy one
+            "labels": torch.zeros_like(input_tensor)
+        }
+
+        # Run inference
+        with torch.no_grad():
+            carry = model.initial_carry(batch)
+            # The model runs in a loop; for inference, we run it for the max steps
+            for _ in range(model_config['halt_max_steps']):
+                carry, _ = model(carry=carry, batch=batch)
+            
+            # Get the final logits from the carry state
+            final_logits = model.inner.lm_head(carry.inner_carry.z_H)[:, model.inner.puzzle_emb_len:]
+            predicted_ids = torch.argmax(final_logits, dim=-1).squeeze().tolist()
         
-        if config is not None:
-            model_info = test_model_info(config, checkpoint)
-            gr.Markdown(model_info)
-    
-    with gr.Tab("Test Interface"):
-        gr.Markdown("## Abstract Optimization Demo")
-        gr.Markdown("*Note: This is a demonstration interface. Full inference requires integration with the training pipeline.*")
+        # Decode the output
+        optimized_text = tokenizer.decode(predicted_ids, skip_special_tokens=True)
         
-        with gr.Row():
-            with gr.Column():
-                draft_input = gr.Textbox(
-                    label="Draft Abstract",
-                    placeholder="Enter your sub-optimal grant abstract here...",
-                    lines=8,
-                    value="Our study will investigate protein interactions in cancer cells. We believe this research could be important for understanding disease mechanisms."
-                )
-                grant_type = gr.Dropdown(
-                    choices=["R01", "F32", "K99", "R21", "R15"],
-                    label="Grant Type",
-                    value="R01"
-                )
-                optimize_btn = gr.Button("Optimize Abstract", variant="primary")
+        return optimized_text
+
+    except Exception as e:
+        print(f"An error occurred during inference: {e}")
+        return f"Error: Could not process the abstract. Details: {e}"
+
+# --- 5. Create the Gradio Interface ---
+grant_type_choices = list(activity_code_map.keys())
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🚀 HRM Grant Abstract Optimizer")
+    gr.Markdown("Enter a draft abstract and select the grant type to get a version optimized by the fine-tuned Hierarchical Reasoning Model.")
+    
+    with gr.Row():
+        with gr.Column():
+            draft_input = gr.Textbox(label="Draft Abstract", lines=15, placeholder="Paste your draft abstract here...")
+            grant_type = gr.Dropdown(label="Grant Type", choices=grant_type_choices, value=grant_type_choices[0] if grant_type_choices else None)
+            optimize_btn = gr.Button("Optimize Abstract", variant="primary")
+        with gr.Column():
+            output_text = gr.Textbox(label="Optimized Abstract", lines=17, interactive=False)
             
-            with gr.Column():
-                output = gr.Textbox(
-                    label="Optimized Abstract",
-                    lines=10,
-                    interactive=False
-                )
-        
-        optimize_btn.click(
-            fn=placeholder_inference,
-            inputs=[draft_input, grant_type],
-            outputs=output
-        )
+    optimize_btn.click(
+        fn=optimize_abstract,
+        inputs=[draft_input, grant_type],
+        outputs=output_text
+    )
 
 if __name__ == "__main__":
     demo.launch()