InternLM2-7B-Recompensa

Proporcionamos algunas API fáciles de usar para que utilices el modelo. Aquí hay un ejemplo de cómo usar el modelo para obtener la puntuación de recompensa de un chat, comparar dos chats o clasificar múltiples chats.

import torch
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained(
"internlm/internlm2-7b-recompensa", 
dispositivo_map="cuda", 
torch_dtype=torch.float16, 
trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained("internlm/internlm2-7b-recompensa", trust_remote_code=True)

chat_1 = [
{"rol": "usuario", "contenido": "¡Hola! ¿Cómo te llamas?"},
{"rol": "asistente", "contenido": "¡Me llamo InternLM2! Un asistente de IA útil. ¿Qué puedo hacer por ti?"}
]
chat_2 = [
{"rol": "usuario", "contenido": "¡Hola! ¿Cómo te llamas?"}, 
{"rol": "asistente", "contenido": "No tengo idea."}
]

# Obtener puntuaciones de recompensa para un solo chat
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 en lote, 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 chats, devolverá el índice de clasificación de cada chat
# el chat 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) # índice más bajo significa puntuación más alta
# >>> rank_res:  [0, 1]

# Ejemplo de uso del modelo de recompensa para realizar Best of N sampling.
# Código que muestra cómo seleccionar la mejor respuesta de los candidatos
generados por el modelo de lenguaje.
import torch
from transformers import AutoModel, AutoTokenizer

# preparar el modelo llm y el tokenizer
llm = AutoModel.from_pretrained(
"internlm/internlm2-chat-7b",
device_map="cuda", 
torch_dtype=torch.float16, 
trust_remote_code=True,
)
llm_tokenizer = AutoTokenizer.from_pretrained("internlm/internlm2-chat-7b", trust_remote_code=True)

# preparar el modelo de recompensa y el tokenizer
reward = AutoModel.from_pretrained(
"internlm/internlm2-7b-recompensa",
device_map="cuda", 
torch_dtype=torch.float16, 
trust_remote_code=True,
)
reward_tokenizer = AutoTokenizer.from_pretrained("internlm/internlm2-7b-recompensa", trust_remote_code=True)

# preparar el prompt de chat
prompt = "Escribe un artículo sobre la revolución de la inteligencia artificial."
messages = [
{"rol": "sistema", "contenido": "Eres un asistente útil."},
{"rol": "usuario", "contenido": prompt}
]
text = llm_tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
model_inputs = llm_tokenizer([text], return_tensors="pt").to("cuda")

# generar N candidatos
num_candidates = 10 # N=10
candidates = []

outputs = llm.generate(
**model_inputs,
max_new_tokens=512,
num_return_sequences=num_candidates,
pad_token_id=llm_tokenizer.eos_token_id,
do_sample=True,
top_k=50,
top_p=0.95,
temperature=0.8,
)
outputs = outputs[:, model_inputs["input_ids"].shape[1]:]

for i in range(num_candidates):
candidate = llm_tokenizer.decode(outputs[i], skip_special_tokens=True)
candidates.append(messages + [{"rol": "asistente", "contenido": candidate}])

rank_indices = reward.rank(reward_tokenizer, candidates)
sorted_candidates = sorted(zip(rank_indices, candidates), key=lambda x: x[0])

## imprimir los candidatos ordenados
# for i, (rank_index, candidate) in enumerate(sorted_candidates):
#     print(f"------------Rank {i}------------: 
{candidate[-1]['content']}")

# imprimir la mejor respuesta
best_response = sorted_candidates[0][1][-1]["contenido"]
print(best_response)