Tarive commited on
Commit
5bb6ee3
·
verified ·
1 Parent(s): 42fbc29

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -25
app.py CHANGED
@@ -5,80 +5,61 @@ import json
5
  from tokenizers import Tokenizer
6
 
7
  # --- 1. Load Custom Model Code ---
8
- # This dynamically loads your corrected HRM source code.
9
- # Since the files are in the root, we import directly.
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, # You may need to get this from your dataset metadata
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() # Set the model to evaluation mode
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) # Default to 0 if unknown
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()