InternLM2-1.8B-Recompensa

import torch
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained(
"internlm/internlm2-1_8b-reward", 
device_map="cuda", 
torch_dtype=torch.float16, 
trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained("internlm/internlm2-1_8b-reward", trust_remote_code=True)

chat_1 = [
{"role": "user", "content": "Hola, ¿Cómo te llamas?"},
{"role": "assistant", "content": "¡Me llamo InternLM2! Un asistente de IA útil. ¿En qué puedo ayudarte?"}
]
chat_2 = [
{"role": "user", "content": "Hola, ¿Cómo te llamas?"}, 
{"role": "assistant", "content": "No tengo idea."}
]

# Obtener puntuación de recompensa para una única conversación
score1 = model.get_score(tokenizer, chat_1)
score2 = model.get_score(tokenizer, chat_2)
print("score1: ", score1)
print("score2: ", score2)
# >>> score1:  0.767578125
# >>> score2:  -2.22265625

# Inferencia por lotes, obtener múltiples puntuaciones a la vez
scores = model.get_scores(tokenizer, [chat_1, chat_2])
print("scores: ", scores)
# >>> scores:  [0.767578125, -2.22265625]

# Comparar si chat_1 es mejor que chat_2
compare_res = model.compare(tokenizer, chat_1, chat_2)
print("compare_res: ", compare_res)
# >>> compare_res:  True

# Clasificar múltiples conversaciones, devolverá el índice de clasificación de cada una
# La conversación con la puntuación más alta tendrá el índice de clasificación como 0
rank_res = model.rank(tokenizer, [chat_1, chat_2])
print("rank_res: ", rank_res) # Un índice más bajo significa una puntuación más alta
# >>> rank_res:  [0, 1]