test

config.json

@@ -18,7 +18,7 @@
   "custom_pipelines": {
     "language-detection": {
       "impl": "impresso_langident_wrapper.Pipeline_One",
-      "pt": ["AutoModel"],
+      "pt": ["DummyModel"],
       "tf": []
     }
   }

impresso_langident_wrapper.py

@@ -5,48 +5,35 @@ from huggingface_hub import hf_hub_download
 
 
 
-class Pipeline_One(Pipeline):
-    # def __init__(self, model_path: str):
-    #     """
-    #     Initialize the Floret language detection pipeline
-        
-    #     Args:
-    #         model_path (str): Path to the .bin model file
-    #     """
-    #     super().__init__()
-    #     self.model = floret.FastText.load_model(model_path)
-    
-    def __init__(self, model_name="floret_model.bin", repo_id="Maslionok/pipeline1", revision="main", **kwargs):
+from transformers import Pipeline
+import floret
+from huggingface_hub import hf_hub_download
+
+class Pipeline_One(Pipeline):  
+    def __init__(self, model=None, model_name="floret_model.bin", repo_id="Maslionok/pipeline1", revision="main", **kwargs):
         """
-        Initialize the Floret language detection pipeline.
-        Args:
-            model_name (str): The name of the Floret model file.
-            repo_id (str): The Hugging Face repository ID.
-            revision (str): The branch/revision to download from.
+        Custom pipeline to manually load a Floret model.
         """
-        super().__init__(**kwargs)  # ✅ Call parent constructor
+        super().__init__(**kwargs)  # ✅ Ensures Hugging Face registers this correctly
 
-        # ✅ Manually download the Floret model
-        model_path = hf_hub_download(repo_id=repo_id, filename=model_name, revision=revision)
-
-        # ✅ Load the Floret model
-        self.model = floret.load_model(model_path)
-
-    def _sanitize_parameters(self, **kwargs):
-        # Add any additional parameter handling if necessary
-        return {}, {}, {}
+        if model is None:
+            # ✅ Manually download the Floret model
+            model_path = hf_hub_download(repo_id=repo_id, filename=model_name, revision=revision)
+            self.model = floret.load_model(model_path)  # ✅ Load Floret manually
+        else:
+            self.model = model  # If manually passed, use existing model
 
     def preprocess(self, text, **kwargs):
-        return text
+        return {"text": text}  # ✅ Prepare text for Floret
 
     def _forward(self, inputs):
-        model_output = self.model.predict(**inputs, k=1)
-
-        return model_output
+        text = inputs["text"]
+        predictions = self.model.predict(text, k=1)  # ✅ Get prediction from Floret
+        return predictions
 
     def postprocess(self, outputs, **kwargs):
-        return outputs
-    
+        return outputs  # ✅ Return Floret’s output as-is
+
 
 
 PIPELINE_REGISTRY.register_pipeline(