Update modelling_magiv2.py

modelling_magiv2.py

@@ -1,8 +1,7 @@
-from transformers import PreTrainedModel, VisionEncoderDecoderModel, ViTMAEModel, ConditionalDetrModel
+from transformers import PreTrainedModel, VisionEncoderDecoderModel, ViTMAEModel, ConditionalDetrModel, ConditionalDetrConfig
 from transformers.models.conditional_detr.modeling_conditional_detr import (
     ConditionalDetrMLPPredictionHead, 
     ConditionalDetrModelOutput,
-    ConditionalDetrHungarianMatcher,
     inverse_sigmoid,
 )
 from .configuration_magiv2 import Magiv2Config
@@ -55,14 +54,14 @@ class Magiv2Model(PreTrainedModel):
             self.class_labels_classifier = nn.Linear(
                 config.detection_model_config.d_model, config.detection_model_config.num_labels
             )
-            self.is_this_text_a_dialogue = nn.Linear(
+                        self.is_this_text_a_dialogue = nn.Linear(
                 config.detection_model_config.d_model, 1
             )
-            self.matcher = ConditionalDetrHungarianMatcher(
+            self.matcher = ConditionalDetrModel(ConditionalDetrConfig(
                 class_cost=config.detection_model_config.class_cost,
                 bbox_cost=config.detection_model_config.bbox_cost,
                 giou_cost=config.detection_model_config.giou_cost
-            )
+            ))
 
     def move_to_device(self, input):
         return move_to_device(input, self.device)