ocr_qa_assessment.py

from transformers import Pipeline


class QAAssessmentPipeline(Pipeline):

    def _sanitize_parameters(self, **kwargs):
        preprocess_kwargs = {}
        if "text" in kwargs:
            preprocess_kwargs["text"] = kwargs["text"]
        return preprocess_kwargs, {}, {}

    def preprocess(self, text, **kwargs):
        # Nothing to preprocess
        return text

    def _forward(self, text, **kwargs):
        predictions, probabilities = self.model(text)
        return predictions, probabilities

    def postprocess(self, outputs, **kwargs):
        predictions = outputs
        print(f"Predictions: {predictions}")
        label = predictions[0][0].replace("__label__", "")  # Remove __label__ prefix

        # Format as JSON-compatible dictionary
        model_output = {"label": label, "score": round(score, 4)}
        return model_output