File size: 1,131 Bytes
dc2b383
db3682f
c746e15
154ad0d
d521794
 
db3682f
26d45f0
c746e15
 
 
154ad0d
 
 
 
 
 
 
c746e15
 
 
 
dc2b383
d521794
86cfb21
dc2b383
d521794
 
dc2b383
26d45f0
c746e15
f202570
d521794
dc2b383
d521794
 
 
1c7044c
d521794
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from transformers import Pipeline
from transformers.pipelines import PIPELINE_REGISTRY
import floret
from huggingface_hub import hf_hub_download



class Pipeline_One(Pipeline):
    def __init__(self, model_path, **kwargs):
        super().__init__(**kwargs)  # Call the base class constructor

        repo_id = "Maslionok/pipeline1"  # Change this to your repo name
        filename = "LID-40-3-2000000-1-4.bin"  # Name of the .bin file
        branch = "main"  # Change this to your branch name if needed

        # Download the model file from Hugging Face
        model_path = hf_hub_download(repo_id=repo_id, filename=filename, revision=branch)

        # Load the Floret model
        self.model = floret.load_model(model_path)  


    def _sanitize_parameters(self, **kwargs):
        # Add any additional parameter handling if necessary
        return {}, {}, {}

    def preprocess(self, text, **kwargs):
        return text

    def _forward(self, inputs):
        model_output = self.model.predict(**inputs, k=1)

        return model_output

    def postprocess(self, outputs, **kwargs):
        return outputs