Dongjin-kr/ko-reranker

Para usar el modelo con Transformers:

def exp_normalize(x):
    b = x.max()
    y = np.exp(x - b)
    return y / y.sum()
    
from transformers import AutoModelForSequenceClassification, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
model.eval()

pairs = [["나는 너를 싫어해", "나는 너를 사랑해"], ["나는 너를 좋아해", "너에 대한 나의 감정은 사랑 일 수도 있어"]]

with torch.no_grad():
    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
    scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
    scores = exp_normalize(scores.numpy())
    print(f'first: {scores[0]}, second: {scores[1]}')

Para usar con SageMaker:

import sagemaker
import boto3
from sagemaker.huggingface import HuggingFaceModel

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

hub = {
    'HF_MODEL_ID': 'Dongjin-kr/ko-reranker',
    'HF_TASK': 'text-classification'
}

huggingface_model = HuggingFaceModel(
    transformers_version='4.28.1',
    pytorch_version='2.0.0',
    py_version='py310',
    env=hub,
    role=role, 
)

predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type='ml.g5.large' 
)

runtime_client = boto3.Session().client('sagemaker-runtime')
payload = json.dumps({"inputs": [
    {"text": "나는 너를 싫어해", "text_pair": "나는 너를 사랑해"},
    {"text": "나는 너를 좋아해", "text_pair": "너에 대한 나의 감정은 사랑 일 수도 있어"}
]})

response = runtime_client.invoke_endpoint(
    EndpointName="",
    ContentType="application/json",
    Accept="application/json",
    Body=payload
)

out = json.loads(response['Body'].read().decode())
print(f'Response: {out}')