Update app.py

app.py

@@ -19,6 +19,22 @@ import gradio as gr
 from entailment_inference import get_scores
 from nle_inference import VideoCaptionDataset, get_nle
 
+import re
+
+def modify_keys(state_dict):
+    new_state_dict = defaultdict()
+
+    pattern = re.compile(r'.*language_model.*\.(q_proj|v_proj|k_proj|o_proj|gate_proj|down_proj|up_proj).weight')
+
+    for key, value in state_dict.items():
+        if pattern.match(key):
+            key = key.split('.')
+            key.insert(-1, 'base_layer')
+            key = '.'.join(key)
+        new_state_dict[key] = value
+
+    return new_state_dict
+
 pretrained_ckpt = "mplugowl7bvideo/"
 trained_ckpt = "owl-con/checkpoint-5178/pytorch_model.bin"
 
@@ -42,8 +58,11 @@ peft_config = LoraConfig(
 )
 model = get_peft_model(model, peft_config)
 model.print_trainable_parameters()
+
 with open(trained_ckpt, 'rb') as f:
     ckpt = torch.load(f, map_location = torch.device("cpu"))
+    ckpt = modify_keys(ckpt)
+    
 model.load_state_dict(ckpt)
 model = model.to("cuda:0").to(torch.bfloat16)