File size: 730 Bytes
dc2b383 db3682f c746e15 d521794 db3682f 26d45f0 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 |
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 {}, {}, {}
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
|