Spaces:
Running
Running
import gradio as gr | |
import torch | |
import yaml | |
import json | |
from tokenizers import Tokenizer | |
from collections import OrderedDict | |
# --- 1. Load Custom Model Code --- | |
# This import now works because we have the correct models/hrm/ structure | |
from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1 | |
# --- 2. Load Artifacts --- | |
print("Loading artifacts...") | |
tokenizer = Tokenizer.from_file("tokenizer.json") | |
with open('config.yaml', 'r') as f: | |
config_data = yaml.safe_load(f) | |
model_config = config_data['arch'] | |
with open('activity_code_map.json', 'r') as f: | |
activity_code_map = json.load(f) | |
# --- 3. Initialize the Model --- | |
print("Initializing model...") | |
model_config.update({ | |
'batch_size': config_data['global_batch_size'], | |
'seq_len': 512, | |
'num_puzzle_identifiers': len(activity_code_map) + 1, | |
'vocab_size': tokenizer.get_vocab_size() | |
}) | |
model = HierarchicalReasoningModel_ACTV1(config_dict=model_config) | |
# --- MODIFICATION: Clean the state dict keys before loading --- | |
# Load the original state dict | |
original_state_dict = torch.load('pytorch_model.bin', map_location='cpu') | |
# Create a new state dict with the corrected keys | |
new_state_dict = OrderedDict() | |
for k, v in original_state_dict.items(): | |
if k.startswith('_orig_mod.model.'): | |
name = k[len('_orig_mod.model.'):] # remove the prefix | |
new_state_dict[name] = v | |
else: | |
new_state_dict[k] = v | |
# Load the cleaned state dict | |
model.load_state_dict(new_state_dict) | |
# --- END MODIFICATION --- | |
model.eval() # Set the model to evaluation mode | |
print("Model loaded successfully!") | |
# --- 4. Define the Inference Function --- | |
def optimize_abstract(draft_abstract, grant_type): | |
if not draft_abstract or not grant_type: | |
return "Please provide both a draft abstract and a grant type." | |
try: | |
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) | |
input_tensor = torch.tensor([input_ids], dtype=torch.long) | |
grant_tensor = torch.tensor([grant_type_id], dtype=torch.long) | |
batch = { | |
"inputs": input_tensor, | |
"puzzle_identifiers": grant_tensor, | |
"labels": torch.zeros_like(input_tensor) | |
} | |
with torch.no_grad(): | |
carry = model.initial_carry(batch) | |
for _ in range(model_config['halt_max_steps']): | |
carry, _ = model(carry=carry, batch=batch) | |
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() | |
optimized_text = tokenizer.decode(predicted_ids, skip_special_tokens=True) | |
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) | |
optimize_btn.click( | |
fn=optimize_abstract, | |
inputs=[draft_input, grant_type], | |
outputs=output_text | |
) | |
if __name__ == "__main__": | |
demo.launch() | |