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, **kwargs):
    #     self.model = model
    #     self.framework = "floret"

    #     super().__init__(model=model)


    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