| 145 | return pred |
| 146 | |
| 147 | class LlamaWithLoss(nn.Module): |
| 148 | def __init__(self, llama, predictor): |
| 149 | super(LlamaWithLoss, self).__init__() |
| 150 | self.llama = llama |
| 151 | self.loss_func = torch.nn.MSELoss() |
| 152 | |
| 153 | self.predictor = predictor |
| 154 | |
| 155 | def forward(self, inputs, y, pooling_method='last_token', return_loss=True): |
| 156 | outputs = self.llama(**inputs, output_hidden_states=True) |
| 157 | last_hidden_state = outputs.last_hidden_state |
| 158 | if pooling_method == 'mean': |
| 159 | embeddings = last_hidden_state.mean(dim=1) # Mean pooling to get sentence-level embeddings |
| 160 | elif pooling_method=='last_token': |
| 161 | embeddings = last_hidden_state[:,-1,:] |
| 162 | if return_loss: |
| 163 | pred = self.predictor(embeddings) |
| 164 | loss = self.loss_func(pred.view(-1),y.view(-1)) |
| 165 | return embeddings, loss |
| 166 | else: |
| 167 | pred = self.predictor(embeddings) |
| 168 | return embeddings, pred |
| 169 | |
| 170 | def read_data_from_csv(path): |
| 171 | data_df = pd.read_csv(path) |