Fix to not break on Papercut model, but still not importing it correctly. Debugging statements enabled.

Author: CrudeDiatribe

backends/model_converter/no_pickle_fake_torch.py

@@ -11,6 +11,7 @@
 else:
     from astunparse import unparse
 
+NO_PICKLE_DEBUG = True
 
 def extract_weights_from_checkpoint(fb0):
   torch_weights = {}
@@ -21,11 +22,13 @@ def extract_weights_from_checkpoint(fb0):
       raise ValueError("Looks like the checkpoints file is in the wrong format")
     folder_name = folder_name[0].replace("/data.pkl" , "").replace("\\data.pkl" , "")
     with myzip.open(folder_name+'/data.pkl') as myfile:
-      instructions = examine_pickle(myfile)
-      for sd_key,load_instruction in instructions.items():
+      load_instructions, special_instructions = examine_pickle(myfile)
+      for sd_key,load_instruction in load_instructions.items():
         with myzip.open(folder_name + f'/data/{load_instruction.obj_key}') as myfile:
           if (load_instruction.load_from_file_buffer(myfile)):
             torch_weights['state_dict'][sd_key] = load_instruction.get_data()
+      for sd_key,special in special_instructions.items():
+        torch_weights['state_dict'][sd_key] = special
   return torch_weights
 
 def examine_pickle(fb0):
@@ -47,6 +50,7 @@ def examine_pickle(fb0):
   re_rebuild = re.compile('^_var\d+ = _rebuild_tensor_v2\(UNPICKLER\.persistent_load\(\(.*\)$')
   re_assign = re.compile('^_var\d+ = \{.*\}$')
   re_update = re.compile('^_var\d+\.update\(\{.*\}\)$')
+  re_ordered_dict = re.compile('^_var\d+ = OrderedDict\(\)$')
 
   load_instructions = {}
   assign_instructions = AssignInstructions()
@@ -61,18 +65,25 @@ def examine_pickle(fb0):
       assign_instructions.parse_assign_line(line)
     elif re_update.match(line):
       assign_instructions.parse_update_line(line)
-    #else:
-    #  print('kicking rocks')
+    elif re_ordered_dict.match(line):
+      #do nothing
+      continue
+    else:
+      print(f'unmatched line: {line}')
+
 
-  #print(f"Found {len(load_instructions)} load instructions")
+  if NO_PICKLE_DEBUG:
+    print(f"Found {len(load_instructions)} load instructions")
 
   assign_instructions.integrate(load_instructions)
 
-  return assign_instructions.integrated_instructions
+  #return assign_instructions.integrated_instructions, assign_instructions.special_instructions
+  return assign_instructions.integrated_instructions, {}
 
 class AssignInstructions:
   def __init__(self):
     self.instructions = {}
+    self.special_instructions = {}
     self.integrated_instructions = {}
 
   def parse_assign_line(self, line):
@@ -84,20 +95,53 @@ def parse_assign_line(self, line):
     assignments = huge_mess.split(', ')
     del huge_mess
     assignments[-1] = assignments[-1].strip('}')
+    re_var = re.compile('^_var\d+$')
+    assignment_count
     for a in assignments:
-      self._add_assignment(a)
-    #print(f"Added/merged {len(assignments)} assignments. Total of {len(self.instructions)} assignment instructions")
-
-  def _add_assignment(self, assignment):
-    sd_key, fickling_var = assignment.split(': ')
+      if self._add_assignment(a, re_var):
+        assignment_count = assignment_count + 1
+    if NO_PICKLE_DEBUG:
+      print(f"Added/merged {assignment_count} assignments. Total of {len(self.instructions)} assignment instructions")
+
+  def _add_assignment(self, assignment, re_var):
+    # assignment can look like this:
+    # 'cond_stage_model.transformer.text_model.encoder.layers.3.self_attn.out_proj.weight': _var2009
+    # or assignment can look like this:
+    # 'embedding_manager.embedder.transformer.text_model.encoder.layers.6.mlp.fc1': {'version': 1}
+    sd_key, fickling_var = assignment.split(': ', 1)
     sd_key = sd_key.strip("'")
-    self.instructions[sd_key] = fickling_var
+    if re_var.match(fickling_var):
+      self.instructions[sd_key] = fickling_var
+      return True
+    else:
+      # now convert the string "{'version': 1}" into a dictionary {'version': 1}
+      entries = fickling_var.split(',')
+      special_dict = {}
+      for e in entries:
+        e = e.strip("{}")
+        k, v = e.split(': ')
+        k = k.strip("'")
+        v = v.strip("'")
+        special_dict[k] = v
+      self.special_instructions[sd_key] = special_dict
+      return False
 
   def integrate(self, load_instructions):
+    unfound_keys = {}
     for sd_key, fickling_var in self.instructions.items():
       if fickling_var in load_instructions:
         self.integrated_instructions[sd_key] = load_instructions[fickling_var]
-    #print(f"Have {len(self.integrated_instructions)} integrated load/assignment instructions")
+      if sd_key in self.special_instructions:
+        if NO_PICKLE_DEBUG:
+          print(f"Key found in both load and special instructions: {sd_key}")
+      else:
+        unfound_keys[sd_key] = True;
+    #for sd_key, special in self.special_instructions.items():
+    #  if sd_key in unfound_keys:
+    #    #todo
+
+    if NO_PICKLE_DEBUG:
+      print(f"Have {len(self.integrated_instructions)} integrated load/assignment instructions")
 
   def parse_update_line(self, line):
     # input looks like:
@@ -109,9 +153,13 @@ def parse_update_line(self, line):
     updates = huge_mess.split(', ')
     del huge_mess
     updates[-1] = updates[-1].strip('})')
+    re_var = re.compile('^_var\d+$')
+    update_count = 0
     for u in updates:
-      self._add_assignment(u)
-    #print(f"Added/merged {len(updates)} updates. Total of {len(self.instructions)} assignment instructions")
+      if self._add_assignment(u, re_var):
+        update_count = update_count + 1
+    if NO_PICKLE_DEBUG:
+      print(f"Added/merged {update_count} updates. Total of {len(self.instructions)} assignment instructions")
 
 class LoadInstruction:
   def __init__(self, instruction_string):