loading model

impresso_langident_wrapper.py

@@ -1,9 +1,17 @@
 from transformers import Pipeline
 from transformers.pipelines import PIPELINE_REGISTRY
+import floret
 
 
 
 class Pipeline_One(Pipeline):
+    def __init__(self, model_path, **kwargs):
+        super().__init__(**kwargs)  # Call the base class constructor
+
+        # Load the Floret model
+        self.model = floret.load_model(model_path)  
+
+
     def _sanitize_parameters(self, **kwargs):
         # Add any additional parameter handling if necessary
         return {}, {}, {}
@@ -12,7 +20,7 @@ class Pipeline_One(Pipeline):
         return text
 
     def _forward(self, inputs):
-        model_output = self.model(**inputs, k=1)
+        model_output = self.model.predict(**inputs, k=1)
 
         return model_output