File size: 659 Bytes
cd6775d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | from sentence_transformers.cross_encoder import CrossEncoder as CE
import numpy as np
from typing import List, Dict, Tuple
class CrossEncoder:
def __init__(self, model_path: str = None, max_length: int = None, **kwargs):
if max_length != None:
self.model = CE(model_path, max_length=max_length, **kwargs)
self.model = CE(model_path, **kwargs)
def predict(self, sentences: List[Tuple[str,str]], batch_size: int = 32, show_progress_bar: bool = True) -> List[float]:
return self.model.predict(
sentences=sentences,
batch_size=batch_size,
show_progress_bar=show_progress_bar) |