|
|
|
|
|
from transformers import PreTrainedModel, Pipeline |
|
from typing import Any, Dict |
|
from transformers import Pipeline |
|
from transformers import PreTrainedTokenizer |
|
from transformers.utils import ModelOutput |
|
|
|
from transformers import PreTrainedModel, Pipeline |
|
from typing import Any, Dict, List |
|
|
|
class QApipeline(Pipeline): |
|
def __init__( |
|
self, |
|
model: PreTrainedModel, |
|
tokenizer: PreTrainedTokenizer, |
|
**kwargs |
|
): |
|
super().__init__( |
|
model=model, |
|
**kwargs |
|
) |
|
|
|
print("in __init__") |
|
|
|
def __call__( self, inputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: |
|
outputs = self.model.predict(inputs) |
|
answer = self._process_output(outputs) |
|
print("in __call___") |
|
return answer |
|
|
|
def _process_output(self, outputs: Any) -> str: |
|
print("in process outputs") |
|
format = {'guess': outputs[1], 'confidence': outputs[0]} |
|
return format |
|
|
|
def _sanitize_parameters(self, **kwargs): |
|
print("in sanitize params") |
|
return {}, {}, {} |
|
|
|
def preprocess(self, inputs): |
|
print("in preprocess") |
|
return inputs |
|
|
|
def postprocess(self, outputs): |
|
print("in postprocess") |
|
format = {'guess': outputs[1], 'confidence': float(outputs[0])} |
|
return format |
|
|
|
def _forward(self, input_tensors, **forward_parameters: Dict) -> ModelOutput: |
|
print("in _forward") |
|
return super()._forward(input_tensors, **forward_parameters) |
|
|
