Update modeling_vit5_kg.py

modeling_vit5_kg.py

@@ -293,14 +293,66 @@ class KGEnhancedViT5(T5ForConditionalGeneration):
         )
         
         return generated_ids
-    
+        
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
         """
         Load model from pretrained path
         """
-        return super().from_pretrained(
-            pretrained_model_name_or_path,
-            *model_args,
-            **kwargs
-        )
+        from transformers import AutoConfig
+        from huggingface_hub import hf_hub_download
+        
+        # Load config
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        
+        # Extract custom fields
+        use_kg = getattr(config, 'use_kg', True)
+        kg_node_features = getattr(config, 'kg_node_features', 300)
+        gnn_hidden = getattr(config, 'gnn_hidden', 256)
+        gnn_type = getattr(config, 'gnn_type', 'gcn')
+        gnn_layers = getattr(config, 'gnn_layers', 2)
+        dropout = getattr(config, 'dropout_rate', 0.1)
+        
+        # Create model với đúng architecture
+        model = cls(
+            config=config,
+            kg_node_features=kg_node_features,
+            gnn_hidden=gnn_hidden,
+            gnn_type=gnn_type,
+            gnn_layers=gnn_layers,
+            dropout=dropout,
+            use_kg=use_kg
+        )
+        
+        # Load weights trực tiếp từ checkpoint file
+        try:
+            # Try safetensors first
+            try:
+                state_dict_path = hf_hub_download(
+                    repo_id=pretrained_model_name_or_path,
+                    filename="model.safetensors"
+                )
+                from safetensors.torch import load_file
+                state_dict = load_file(state_dict_path)
+            except:
+                # Fallback to pytorch_model.bin
+                state_dict_path = hf_hub_download(
+                    repo_id=pretrained_model_name_or_path,
+                    filename="pytorch_model.bin"
+                )
+                state_dict = torch.load(state_dict_path, map_location='cpu')
+            
+            # Load weights với strict=False
+            model.load_state_dict(state_dict, strict=False)
+        except Exception as e:
+            # Fallback: dùng parent's from_pretrained
+            print(f"Warning: Could not load weights directly: {e}")
+            parent_model = super().from_pretrained(
+                pretrained_model_name_or_path,
+                *model_args,
+                config=config,
+                **kwargs
+            )
+            model.load_state_dict(parent_model.state_dict(), strict=False)
+        
+        return model