Views
No views yet
1class P2LOutputs(ModelOutput):
2 coefs: torch.FloatTensor = None # "betas" as described above
3 eta: Optional[torch.FloatTensor] = None # tie coefficent (not used for BT head)
4 last_hidden_state: torch.FloatTensor = None # last hidden state from the transformermodel_list.json found in the repo of each P2L model. As a general rule, the models will always be in sorted order.1import json
2from huggingface_hub import hf_hub_download
3
4fname = hf_hub_download(
5 repo_id="lmarena-ai/p2l-360m-bt-01132025", filename="model_list.json", repo_type="model"
6 )
7
8with open(fname) as fin:
9 model_list = json.load(fin)1
2import torch
3from transformers import (
4 Qwen2Model,
5 Qwen2PreTrainedModel,
6 LlamaModel,
7 LlamaPreTrainedModel,
8 PreTrainedModel,
9 AutoTokenizer,
10)
11from transformers import AutoTokenizer
12from transformers.utils import ModelOutput
13from dataclasses import dataclass
14import torch.nn as nn
15import torch.nn.functional as F
16from typing import Dict, Tuple, Callable, Optional
17from huggingface_hub import hf_hub_download
18import json
19
20
21@dataclass
22class HeadOutputs(ModelOutput):
23 coefs: torch.FloatTensor = None
24 eta: Optional[torch.FloatTensor] = None
25 gamma: Optional[torch.FloatTensor] = None
26
27
28@dataclass
29class P2LOutputs(ModelOutput):
30 coefs: torch.FloatTensor = None
31 eta: Optional[torch.FloatTensor] = None
32 gamma: Optional[torch.FloatTensor] = None
33 loss: Optional[torch.FloatTensor] = None
34 last_hidden_state: torch.FloatTensor = None
35
36class BTHead(nn.Module):
37 def __init__(
38 self, input_dim, output_dim, linear_head_downsize_factor=None, **kwargs
39 ) -> None:
40 super().__init__()
41
42 if linear_head_downsize_factor:
43 inner_dim = int(output_dim // linear_head_downsize_factor)
44 self.head = nn.Sequential(
45 nn.Linear(in_features=input_dim, out_features=inner_dim, bias=True),
46 nn.Linear(in_features=inner_dim, out_features=output_dim, bias=True),
47 )
48 else:
49 self.head = nn.Linear(
50 in_features=input_dim, out_features=output_dim, bias=True
51 )
52
53 def forward(self, last_hidden_dim: torch.Tensor):
54 coefs = self.head(last_hidden_dim)
55 return HeadOutputs(coefs=coefs)
56
57class P2LModel(LlamaPreTrainedModel):
58 def __init__(
59 self,
60 config,
61 CLS_id,
62 num_models,
63 head_kwargs={},
64 **kwargs,
65 ):
66 super().__init__(config)
67
68 self.num_models = num_models
69 self.cls_token_id = CLS_id
70
71 self.model = LlamaModel(config)
72
73 self.head = BTHead(
74 input_dim=config.hidden_size,
75 output_dim=self.num_models,
76 **head_kwargs,
77 )
78
79 self.post_init()
80
81 def freeze_transformer(self):
82 for param in self.model.parameters():
83 param.requires_grad = False
84
85 def get_input_embeddings(self):
86 return self.model.embed_tokens
87
88 def set_input_embeddings(self, value):
89 self.model.embed_tokens = value
90
91 def forward(self, input_ids, attention_mask, labels=None, weights=None):
92 batch_size = input_ids.shape[0]
93
94 hidden_outputs = self.model(
95 input_ids=input_ids,
96 attention_mask=attention_mask,
97 output_hidden_states=False,
98 ).last_hidden_state # (bs, num_token, embed_dim)
99
100 cls_mask = input_ids == self.cls_token_id
101
102 # double check this is getting the current CLS token
103 cls_hidden_dim = hidden_outputs[cls_mask]
104
105 assert (
106 cls_hidden_dim.shape[0] == batch_size
107 ), f"input ids {input_ids.shape}, cls_mask {cls_mask.shape}, cls_logit {cls_hidden_dim.shape}"
108
109 head_output = self.head(cls_hidden_dim)
110
111
112 outputs = P2LOutputs(
113 coefs=head_output.coefs,
114 last_hidden_state=cls_hidden_dim,
115 eta=head_output.eta,
116 gamma=head_output.gamma,
117 )
118
119 return outputs
120
121
122fname = hf_hub_download(
123 repo_id="lmarena-ai/p2l-360m-bt-01132025", filename="model_list.json", repo_type="model"
124 )
125
126with open(fname) as fin:
127 model_list = json.load(fin)
128
129tokenizer = AutoTokenizer.from_pretrained("lmarena-ai/p2l-360m-bt-01132025")
130model = P2LModel.from_pretrained(
131 "lmarena-ai/p2l-360m-bt-01132025",
132 CLS_id=tokenizer.cls_token_id,
133 num_models=len(model_list),
134 torch_dtype=torch.bfloat16,
135)
136@misc{frick2025prompttoleaderboard,
title={Prompt-to-Leaderboard},
author={Evan Frick and Connor Chen and Joseph Tennyson and Tianle Li and Wei-Lin Chiang and Anastasios N. Angelopoulos and Ion Stoica},
year={2025},
eprint={2502.14855},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2502.14855},
}