fixing connection issue

model.py

@@ -78,16 +78,10 @@ class BertClassifier(LabelStudioMLBase):
         
         # Initialize basic attributes
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        self.version = 'v0.0.1'  # Define version explicitly
-        self.model_dir = f'BertClassifier-{self.version}'  # Use versioned model directory
+        self.version = 'v0.0.1'
+        self.model_dir = f'BertClassifier-{self.version}'
         
-        # Initialize Label Studio client
-        self.label_studio_client = self.connect_to_label_studio()
-        
-        self.label_encoder = LabelEncoder()
-        self.instruction_template = os.getenv('MODEL_INSTRUCTIONS', '{text}')
-        
-        # Define your categories
+        # Define categories
         self.categories = [
             'affiliate_classification', 'brand', 'business_and_career',
             'content_quality', 'date', 'demographic', 'event',
@@ -98,8 +92,11 @@ class BertClassifier(LabelStudioMLBase):
             'style_and_fashion', 'no_category'
         ]
         
-        # Fit label encoder with your categories
-        self.label_encoder.fit(self.categories)
+        # Initialize model and tokenizer as None - they'll be loaded when needed
+        self._model = None
+        self.tokenizer = None
+        
+        logger.info("BertClassifier initialized successfully")
 
     def get_labels(self):
         li = self.label_interface
@@ -108,7 +105,31 @@ class BertClassifier(LabelStudioMLBase):
         return tag.labels
 
     def setup(self):
-        self.set("model_version", f'{self.__class__.__name__}-v0.0.1')
+        """Setup the model - this is called when Label Studio connects"""
+        try:
+            # Initialize model directory
+            os.makedirs(self.model_dir, exist_ok=True)
+            
+            # Return the required information for Label Studio
+            return {
+                'model_class': 'BertClassifier',  # Must match your class name
+                'model_params': {
+                    'device': str(self.device),
+                    'version': self.version
+                },
+                'label_config': {
+                    'from_name': 'sentiment',
+                    'to_name': 'text',
+                    'type': 'choices',
+                    'labels': self.categories
+                },
+                'api_version': '2'  # Important: specify API version
+            }
+            
+        except Exception as e:
+            logger.error(f"Error in setup: {str(e)}")
+            logger.error("Full error details:", exc_info=True)
+            raise
 
     def _lazy_init(self):
         if not hasattr(self, '_model') or self._model is None: