Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -5,80 +5,61 @@ import json
|
|
5 |
from tokenizers import Tokenizer
|
6 |
|
7 |
# --- 1. Load Custom Model Code ---
|
8 |
-
# This
|
9 |
-
|
10 |
-
from hrm_act_v1 import HierarchicalReasoningModel_ACTV1
|
11 |
|
12 |
# --- 2. Load Artifacts ---
|
13 |
print("Loading artifacts...")
|
14 |
-
# Load the tokenizer
|
15 |
tokenizer = Tokenizer.from_file("tokenizer.json")
|
16 |
-
# Load the model configuration
|
17 |
with open('config.yaml', 'r') as f:
|
18 |
config_data = yaml.safe_load(f)
|
19 |
model_config = config_data['arch']
|
20 |
-
# Load the grant type mapping
|
21 |
with open('activity_code_map.json', 'r') as f:
|
22 |
activity_code_map = json.load(f)
|
23 |
|
24 |
# --- 3. Initialize the Model ---
|
25 |
print("Initializing model...")
|
26 |
-
# The model expects a dict, so we pass the Pydantic model's dict representation
|
27 |
-
# We also need to add other required keys from the root of the config
|
28 |
model_config.update({
|
29 |
'batch_size': config_data['global_batch_size'],
|
30 |
-
'seq_len': 512,
|
31 |
'num_puzzle_identifiers': len(activity_code_map) + 1,
|
32 |
'vocab_size': tokenizer.get_vocab_size()
|
33 |
})
|
34 |
model = HierarchicalReasoningModel_ACTV1(config_dict=model_config)
|
35 |
-
# Load the fine-tuned weights
|
36 |
model.load_state_dict(torch.load('pytorch_model.bin', map_location='cpu'))
|
37 |
-
model.eval()
|
38 |
print("Model loaded successfully!")
|
39 |
|
40 |
# --- 4. Define the Inference Function ---
|
41 |
def optimize_abstract(draft_abstract, grant_type):
|
42 |
-
"""
|
43 |
-
Takes a draft abstract and grant type, runs the model, and returns the optimized text.
|
44 |
-
"""
|
45 |
if not draft_abstract or not grant_type:
|
46 |
return "Please provide both a draft abstract and a grant type."
|
47 |
|
48 |
try:
|
49 |
-
# Prepare inputs
|
50 |
tokenizer.enable_padding(length=512)
|
51 |
tokenizer.enable_truncation(max_length=512)
|
52 |
|
53 |
input_ids = tokenizer.encode(draft_abstract).ids
|
54 |
-
grant_type_id = activity_code_map.get(grant_type, 0)
|
55 |
|
56 |
-
# Convert to PyTorch tensors
|
57 |
input_tensor = torch.tensor([input_ids], dtype=torch.long)
|
58 |
grant_tensor = torch.tensor([grant_type_id], dtype=torch.long)
|
59 |
|
60 |
-
# Create the batch dictionary that the model expects
|
61 |
batch = {
|
62 |
"inputs": input_tensor,
|
63 |
"puzzle_identifiers": grant_tensor,
|
64 |
-
# The model requires a 'labels' field, even for inference, so we provide a dummy one
|
65 |
"labels": torch.zeros_like(input_tensor)
|
66 |
}
|
67 |
|
68 |
-
# Run inference
|
69 |
with torch.no_grad():
|
70 |
carry = model.initial_carry(batch)
|
71 |
-
# The model runs in a loop; for inference, we run it for the max steps
|
72 |
for _ in range(model_config['halt_max_steps']):
|
73 |
carry, _ = model(carry=carry, batch=batch)
|
74 |
|
75 |
-
# Get the final logits from the carry state
|
76 |
final_logits = model.inner.lm_head(carry.inner_carry.z_H)[:, model.inner.puzzle_emb_len:]
|
77 |
predicted_ids = torch.argmax(final_logits, dim=-1).squeeze().tolist()
|
78 |
|
79 |
-
# Decode the output
|
80 |
optimized_text = tokenizer.decode(predicted_ids, skip_special_tokens=True)
|
81 |
-
|
82 |
return optimized_text
|
83 |
|
84 |
except Exception as e:
|
@@ -107,4 +88,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
|
|
107 |
)
|
108 |
|
109 |
if __name__ == "__main__":
|
110 |
-
demo.launch()
|
|
|
5 |
from tokenizers import Tokenizer
|
6 |
|
7 |
# --- 1. Load Custom Model Code ---
|
8 |
+
# This import now works because we have the correct models/hrm/ structure
|
9 |
+
from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1
|
|
|
10 |
|
11 |
# --- 2. Load Artifacts ---
|
12 |
print("Loading artifacts...")
|
|
|
13 |
tokenizer = Tokenizer.from_file("tokenizer.json")
|
|
|
14 |
with open('config.yaml', 'r') as f:
|
15 |
config_data = yaml.safe_load(f)
|
16 |
model_config = config_data['arch']
|
|
|
17 |
with open('activity_code_map.json', 'r') as f:
|
18 |
activity_code_map = json.load(f)
|
19 |
|
20 |
# --- 3. Initialize the Model ---
|
21 |
print("Initializing model...")
|
|
|
|
|
22 |
model_config.update({
|
23 |
'batch_size': config_data['global_batch_size'],
|
24 |
+
'seq_len': 512,
|
25 |
'num_puzzle_identifiers': len(activity_code_map) + 1,
|
26 |
'vocab_size': tokenizer.get_vocab_size()
|
27 |
})
|
28 |
model = HierarchicalReasoningModel_ACTV1(config_dict=model_config)
|
|
|
29 |
model.load_state_dict(torch.load('pytorch_model.bin', map_location='cpu'))
|
30 |
+
model.eval()
|
31 |
print("Model loaded successfully!")
|
32 |
|
33 |
# --- 4. Define the Inference Function ---
|
34 |
def optimize_abstract(draft_abstract, grant_type):
|
|
|
|
|
|
|
35 |
if not draft_abstract or not grant_type:
|
36 |
return "Please provide both a draft abstract and a grant type."
|
37 |
|
38 |
try:
|
|
|
39 |
tokenizer.enable_padding(length=512)
|
40 |
tokenizer.enable_truncation(max_length=512)
|
41 |
|
42 |
input_ids = tokenizer.encode(draft_abstract).ids
|
43 |
+
grant_type_id = activity_code_map.get(grant_type, 0)
|
44 |
|
|
|
45 |
input_tensor = torch.tensor([input_ids], dtype=torch.long)
|
46 |
grant_tensor = torch.tensor([grant_type_id], dtype=torch.long)
|
47 |
|
|
|
48 |
batch = {
|
49 |
"inputs": input_tensor,
|
50 |
"puzzle_identifiers": grant_tensor,
|
|
|
51 |
"labels": torch.zeros_like(input_tensor)
|
52 |
}
|
53 |
|
|
|
54 |
with torch.no_grad():
|
55 |
carry = model.initial_carry(batch)
|
|
|
56 |
for _ in range(model_config['halt_max_steps']):
|
57 |
carry, _ = model(carry=carry, batch=batch)
|
58 |
|
|
|
59 |
final_logits = model.inner.lm_head(carry.inner_carry.z_H)[:, model.inner.puzzle_emb_len:]
|
60 |
predicted_ids = torch.argmax(final_logits, dim=-1).squeeze().tolist()
|
61 |
|
|
|
62 |
optimized_text = tokenizer.decode(predicted_ids, skip_special_tokens=True)
|
|
|
63 |
return optimized_text
|
64 |
|
65 |
except Exception as e:
|
|
|
88 |
)
|
89 |
|
90 |
if __name__ == "__main__":
|
91 |
+
demo.launch()
|