dgaff commited on
Commit
5ff507b
·
1 Parent(s): 08cb742

init model

Browse files
Files changed (8) hide show
  1. LICENSE +21 -0
  2. data_processing.py +55 -0
  3. inference.py +48 -0
  4. label_mappings.json +1 -0
  5. multioutput_regressor.pth +3 -0
  6. requirements.txt +7 -0
  7. train.py +301 -0
  8. utils.py +15 -0
LICENSE ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MIT License
2
+
3
+ Copyright (c) 2024 Devin Gaffney
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
data_processing.py ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # data_processing.py
2
+ import json
3
+ import torch
4
+ from transformers import DistilBertTokenizerFast, DistilBertModel
5
+ import numpy as np
6
+
7
+ def load_data(file_path):
8
+ with open(file_path, 'r') as f:
9
+ dataset = json.load(f)
10
+ outdata = [
11
+ {
12
+ "did": e["user_id"],
13
+ "description": e["description"],
14
+ "label_weights": e["user_categories"]
15
+ }
16
+ for e in dataset
17
+ if e["description"] and e["user_categories"]
18
+ ]
19
+ return outdata
20
+
21
+ def prepare_labels(outdata):
22
+ all_labels = sorted({label for record in outdata for label in record['label_weights'].keys()})
23
+ label2id = {label: i for i, label in enumerate(all_labels)}
24
+ id2label = {i: label for label, i in label2id.items()}
25
+
26
+ y_matrix = np.zeros((len(outdata), len(all_labels)), dtype=float)
27
+ for idx, record in enumerate(outdata):
28
+ for label, weight in record['label_weights'].items():
29
+ y_matrix[idx, label2id[label]] = weight
30
+ return y_matrix, label2id, id2label
31
+
32
+ class EmbeddingGenerator:
33
+ def __init__(self, model_name='distilbert-base-uncased', device=None):
34
+ self.tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
35
+ self.embedding_model = DistilBertModel.from_pretrained(model_name)
36
+ self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
37
+ self.embedding_model.to(self.device)
38
+
39
+ def generate_embeddings(self, descriptions, batch_size=1000):
40
+ all_embeddings = []
41
+ descriptions = [desc for desc in descriptions]
42
+ for i in range(0, len(descriptions), batch_size):
43
+ batch_descriptions = descriptions[i:i + batch_size]
44
+ inputs = self.tokenizer(
45
+ batch_descriptions,
46
+ padding=True,
47
+ truncation=True,
48
+ max_length=128,
49
+ return_tensors="pt"
50
+ ).to(self.device)
51
+ with torch.no_grad():
52
+ outputs = self.embedding_model(**inputs)
53
+ batch_embeddings = outputs.last_hidden_state[:, 0, :].cpu().numpy()
54
+ all_embeddings.append(batch_embeddings)
55
+ return np.vstack(all_embeddings)
inference.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # inference.py
2
+ import torch
3
+ import numpy as np
4
+ import joblib
5
+ import json
6
+ from transformers import DistilBertTokenizerFast, DistilBertModel
7
+
8
+ class Predictor:
9
+ def __init__(self, model_path='xgboost_model.joblib', mappings_path='label_mappings.json', device=None):
10
+ # Load the XGBoost model
11
+ self.model = joblib.load(model_path)
12
+
13
+ # Load label mappings
14
+ with open(mappings_path, 'r') as f:
15
+ mappings = json.load(f)
16
+ self.id2label = {int(k): v for k, v in mappings['id2label'].items()}
17
+
18
+ # Load the tokenizer and embedding model
19
+ self.tokenizer = DistilBertTokenizerFast.from_pretrained("distilbert-base-uncased")
20
+ self.embedding_model = DistilBertModel.from_pretrained("distilbert-base-uncased")
21
+ self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
22
+ self.embedding_model.to(self.device)
23
+
24
+ def generate_embedding(self, text):
25
+ inputs = self.tokenizer(
26
+ [text],
27
+ padding=True,
28
+ truncation=True,
29
+ max_length=128,
30
+ return_tensors="pt"
31
+ ).to(self.device)
32
+ with torch.no_grad():
33
+ outputs = self.embedding_model(**inputs)
34
+ embedding = outputs.last_hidden_state[:, 0, :].cpu().numpy()
35
+ return embedding
36
+
37
+ def predict(self, text):
38
+ embedding = self.generate_embedding(text)
39
+ y_pred = self.model.predict(embedding)
40
+ predictions = {self.id2label[i]: float(y_pred[0][i]) for i in range(len(self.id2label))}
41
+ return predictions
42
+
43
+ # Example usage
44
+ if __name__ == "__main__":
45
+ predictor = Predictor()
46
+ text = "I write about American politics"
47
+ predictions = predictor.predict(text)
48
+ print(predictions)
label_mappings.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"label2id": {"Animals": 0, "Art": 1, "Books": 2, "Comedy": 3, "Comics": 4, "Culture": 5, "Education": 6, "Food": 7, "Journalism": 8, "Movies": 9, "Music": 10, "Nature": 11, "News": 12, "Pets": 13, "Photography": 14, "Politics": 15, "Science": 16, "Software Dev": 17, "Sports": 18, "TV": 19, "Tech": 20, "Video Games": 21, "Writers": 22}, "id2label": {"0": "Animals", "1": "Art", "2": "Books", "3": "Comedy", "4": "Comics", "5": "Culture", "6": "Education", "7": "Food", "8": "Journalism", "9": "Movies", "10": "Music", "11": "Nature", "12": "News", "13": "Pets", "14": "Photography", "15": "Politics", "16": "Science", "17": "Software Dev", "18": "Sports", "19": "TV", "20": "Tech", "21": "Video Games", "22": "Writers"}}
multioutput_regressor.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e1a318d1aaf0f83962c6acb25834ccae74413b7686c2622257f91df42f99781d
3
+ size 72364
requirements.txt ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ numpy
2
+ pandas
3
+ torch
4
+ transformers
5
+ scikit-learn
6
+ xgboost
7
+ joblib
train.py ADDED
@@ -0,0 +1,301 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # train.py
2
+
3
+ import numpy as np
4
+ import torch
5
+ from torch import nn
6
+ from torch.utils.data import DataLoader
7
+ from sklearn.model_selection import KFold
8
+ from transformers import Trainer, TrainingArguments
9
+ from sklearn.metrics import ndcg_score
10
+ import json
11
+
12
+ from data_processing import load_data, EmbeddingGenerator, prepare_labels
13
+ from utils import compute_ndcg
14
+ import numpy as np
15
+ from sklearn.metrics import ndcg_score, mean_squared_error
16
+
17
+ # Generate random predictions based on label distribution
18
+ def generate_random_predictions(y_true):
19
+ return np.random.uniform(y_true.min(), y_true.max(), size=y_true.shape)
20
+
21
+ # Evaluate relative lift
22
+ def calculate_relative_lift(y_true, model_predictions, metric="ndcg"):
23
+ random_predictions = generate_random_predictions(y_true)
24
+
25
+ if metric == "ndcg":
26
+ model_score = ndcg_score([y_true], [model_predictions])
27
+ random_score = ndcg_score([y_true], [random_predictions])
28
+ lift = (model_score - random_score) / random_score
29
+ elif metric == "mse":
30
+ model_score = mean_squared_error(y_true, model_predictions)
31
+ random_score = mean_squared_error(y_true, random_predictions)
32
+ lift = (random_score - model_score) / random_score
33
+ else:
34
+ raise ValueError("Unsupported metric")
35
+
36
+ return lift, model_score, random_score
37
+
38
+ # Define your model architecture
39
+ class MultiOutputRegressor(nn.Module):
40
+ def __init__(self, hidden_size, num_outputs):
41
+ super(MultiOutputRegressor, self).__init__()
42
+ self.regressor_head = nn.Linear(hidden_size, num_outputs)
43
+
44
+ def forward(self, input_ids):
45
+ return self.regressor_head(input_ids)
46
+
47
+ # Dataset class
48
+ class EmbeddingDataset(torch.utils.data.Dataset):
49
+ def __init__(self, embeddings, labels):
50
+ self.embeddings = embeddings
51
+ self.labels = labels
52
+
53
+ def __len__(self):
54
+ return len(self.embeddings)
55
+
56
+ def __getitem__(self, idx):
57
+ return {"input_ids": self.embeddings[idx], "label": self.labels[idx]}
58
+
59
+ # Custom data collator
60
+ class CustomDataCollator:
61
+ def __call__(self, features):
62
+ embeddings = torch.stack([item["input_ids"] for item in features])
63
+ labels = torch.stack([item["label"] for item in features])
64
+ batch_data = {"input_ids": embeddings, "label": labels}
65
+ return batch_data
66
+
67
+ # Custom Trainer
68
+ class CustomTrainer(Trainer):
69
+ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
70
+ input_ids = inputs["input_ids"].to(self.args.device)
71
+ labels = inputs["label"].to(self.args.device)
72
+ outputs = model(input_ids)
73
+ loss_fct = nn.MSELoss()
74
+ loss = loss_fct(outputs, labels)
75
+ return (loss, outputs) if return_outputs else loss
76
+
77
+ def main():
78
+ # Load data
79
+ outdata = load_data("labeled_users.json")
80
+
81
+ # Extract descriptions
82
+ descriptions = [record['description'] for record in outdata]
83
+
84
+ # Generate embeddings
85
+ embedder = EmbeddingGenerator()
86
+ X_embeddings = embedder.generate_embeddings(descriptions)
87
+
88
+ # Prepare labels
89
+ y_matrix, label2id, id2label = prepare_labels(outdata)
90
+
91
+ # Save label mappings for later use
92
+ mappings = {'label2id': label2id, 'id2label': id2label}
93
+ with open('label_mappings.json', 'w') as f:
94
+ json.dump(mappings, f)
95
+
96
+ # K-Fold Cross Validation
97
+ train_embeddings = torch.tensor(X_embeddings, dtype=torch.float)
98
+ train_labels = torch.tensor(y_matrix, dtype=torch.float)
99
+
100
+ # Device configuration
101
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
102
+
103
+ data_collator = CustomDataCollator()
104
+
105
+ n_splits = 5 # Number of folds
106
+ kf = KFold(n_splits=n_splits, shuffle=True, random_state=42)
107
+
108
+ hidden_size = train_embeddings.shape[1]
109
+ num_outputs = train_labels.shape[1]
110
+ fold_ndcg_scores = []
111
+ all_preds = []
112
+
113
+ for fold, (train_index, val_index) in enumerate(kf.split(train_embeddings)):
114
+ print(f"Fold {fold + 1}/{n_splits}")
115
+
116
+ # Split data into training and validation sets
117
+ X_train_fold = train_embeddings[train_index]
118
+ y_train_fold = train_labels[train_index]
119
+ X_val_fold = train_embeddings[val_index]
120
+ y_val_fold = train_labels[val_index]
121
+
122
+ # Create datasets
123
+ train_dataset = EmbeddingDataset(X_train_fold, y_train_fold)
124
+ val_dataset = EmbeddingDataset(X_val_fold, y_val_fold)
125
+
126
+ # Initialize the model
127
+ model = MultiOutputRegressor(hidden_size=hidden_size, num_outputs=num_outputs)
128
+ model.to(device)
129
+
130
+ # Training arguments
131
+ training_args = TrainingArguments(
132
+ output_dir=f"./results_fold_{fold+1}",
133
+ num_train_epochs=10,
134
+ per_device_train_batch_size=64,
135
+ logging_dir=f"./logs_fold_{fold+1}",
136
+ evaluation_strategy="no", # No evaluation during training
137
+ save_strategy="no", # Not saving checkpoints
138
+ disable_tqdm=True, # Disable progress bar
139
+ learning_rate=1e-5,
140
+ weight_decay=0.01, # Apply a small weight decay
141
+ max_grad_norm=1.0 # Clip gradients to 1.0
142
+ )
143
+ # Initialize Trainer
144
+ trainer = CustomTrainer(
145
+ model=model,
146
+ args=training_args,
147
+ train_dataset=train_dataset,
148
+ data_collator=data_collator,
149
+ )
150
+ # Train the model
151
+ trainer.train()
152
+
153
+ # Evaluate the model on the validation set
154
+ val_dataloader = DataLoader(val_dataset, batch_size=8, collate_fn=data_collator)
155
+
156
+ fold_preds = []
157
+ fold_labels = []
158
+
159
+ model.eval()
160
+ with torch.no_grad():
161
+ for batch in val_dataloader:
162
+ input_ids = batch["input_ids"].to(device)
163
+ labels = batch["label"].to(device)
164
+ outputs = model(input_ids)
165
+ fold_preds.append(outputs.cpu().numpy())
166
+ fold_labels.append(labels.cpu().numpy())
167
+
168
+ # Concatenate all predictions and labels for the fold
169
+ y_pred = np.concatenate(fold_preds, axis=0)
170
+ y_true = np.concatenate(fold_labels, axis=0)
171
+
172
+ # Append fold predictions to all_preds
173
+ all_preds.extend(y_pred)
174
+
175
+ # Compute NDCG scores
176
+ all_ndcgs = []
177
+ lifts = []
178
+ for i in range(len(y_true)):
179
+ actual_weights = y_true[i]
180
+ predicted_weights = y_pred[i]
181
+ ndcg = ndcg_score([actual_weights], [predicted_weights])
182
+ lift, model_score, random_score = calculate_relative_lift(actual_weights, predicted_weights, metric="ndcg")
183
+ lifts.append(lift)
184
+ all_ndcgs.append(ndcg)
185
+
186
+ # Average NDCG score for the current fold
187
+ if all_ndcgs:
188
+ avg_ndcg = np.mean(all_ndcgs)
189
+ else:
190
+ avg_ndcg = 0.0 # Handle cases where there are no non-zero weights
191
+ if lifts:
192
+ avg_lift = np.mean(lifts)
193
+ else:
194
+ avg_lift = 0.0 # Handle cases where there are no non-zero weights
195
+ print(f"Average NDCG for fold {fold + 1}: {avg_ndcg:.4f}")
196
+ print(f"Average Lift for fold {fold + 1}: {avg_lift:.4f}")
197
+ fold_ndcg_scores.append(avg_ndcg)
198
+
199
+ # After all folds
200
+ overall_avg_ndcg = np.mean(fold_ndcg_scores)
201
+ print(f"\nOverall Average NDCG across all folds: {overall_avg_ndcg:.4f}")
202
+
203
+ # Store embeddings and predictions in outdata
204
+ for idx, record in enumerate(outdata):
205
+ record['embedding'] = X_embeddings[idx].tolist()
206
+ # Map predictions to labels
207
+ pred = all_preds[idx]
208
+ label_pred_dict = {id2label[i]: float(pred[i]) for i in range(len(pred))}
209
+ record['predictions'] = label_pred_dict
210
+
211
+ # Save enriched data
212
+ with open("enriched_data.json", "w") as f:
213
+ for row in outdata:
214
+ _ = f.write(json.dumps(row) + '\n')
215
+
216
+ # Save full model trained on entire dataset
217
+ # Re-initialize the model
218
+ model = MultiOutputRegressor(hidden_size=hidden_size, num_outputs=num_outputs)
219
+ model.to(device)
220
+
221
+ # Create the dataset with all data
222
+ train_dataset = EmbeddingDataset(train_embeddings, train_labels)
223
+
224
+ # Training arguments
225
+ training_args = TrainingArguments(
226
+ output_dir="./final_model",
227
+ num_train_epochs=10, # Adjust as needed
228
+ per_device_train_batch_size=8,
229
+ logging_dir="./logs_final",
230
+ evaluation_strategy="no",
231
+ save_strategy="no",
232
+ disable_tqdm=False,
233
+ )
234
+
235
+ # Initialize the Trainer
236
+ trainer = CustomTrainer(
237
+ model=model,
238
+ args=training_args,
239
+ train_dataset=train_dataset,
240
+ data_collator=data_collator,
241
+ )
242
+
243
+ # Train the model on the entire dataset
244
+ trainer.train()
245
+
246
+ # Save the model
247
+ model_save_path = 'multioutput_regressor.pth'
248
+ torch.save(model.state_dict(), model_save_path)
249
+ print(f"Model saved to {model_save_path}")
250
+
251
+ # Optional: Demonstrate loading and using the model
252
+ load_and_predict(embedder, hidden_size, num_outputs, device)
253
+
254
+ def load_and_predict(embedder, hidden_size, num_outputs, device):
255
+ """
256
+ Load the saved model and label mappings, make predictions on new data,
257
+ and map the predictions to labels.
258
+ """
259
+ # Load the label mappings
260
+ with open('label_mappings.json', 'r') as f:
261
+ mappings = json.load(f)
262
+ id2label = mappings['id2label']
263
+
264
+ # Load the model
265
+ model_save_path = 'multioutput_regressor.pth'
266
+ loaded_model = MultiOutputRegressor(hidden_size=hidden_size, num_outputs=num_outputs)
267
+ loaded_model.load_state_dict(torch.load(model_save_path, map_location=device))
268
+ loaded_model.to(device)
269
+ loaded_model.eval()
270
+
271
+ # Prepare new data for prediction
272
+ new_sentences = [
273
+ "This is a test sentence.",
274
+ "Another example of a sentence to predict."
275
+ ]
276
+ # Generate embeddings for new sentences
277
+ new_embeddings = embedder.generate_embeddings(new_sentences)
278
+ new_embeddings_tensor = torch.tensor(new_embeddings, dtype=torch.float).to(device)
279
+
280
+ # Make predictions
281
+ with torch.no_grad():
282
+ predictions = loaded_model(new_embeddings_tensor)
283
+ predictions = predictions.cpu().numpy()
284
+
285
+ # Map predictions to labels
286
+ for sentence, pred in zip(new_sentences, predictions):
287
+ label_pred_dict = {id2label[str(i)]: float(pred[i]) for i in range(len(pred))}
288
+ print(f"Sentence: {sentence}")
289
+ print("Predictions:")
290
+ for label, value in label_pred_dict.items():
291
+ print(f" {label}: {value}")
292
+ print()
293
+
294
+ if __name__ == "__main__":
295
+ main()
296
+
297
+ # loaded_model = MultiOutputRegressor(hidden_size=hidden_size, num_outputs=num_outputs)
298
+ # loaded_model.load_state_dict(torch.load(model_save_path, map_location=device))
299
+ # loaded_model.to(device)
300
+ # loaded_model.eval()
301
+
utils.py ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+
3
+ def compute_dcg(relevances):
4
+ relevances = np.asarray(relevances)
5
+ discounts = np.log2(np.arange(len(relevances)) + 2)
6
+ return np.sum(relevances / discounts)
7
+
8
+ def compute_ndcg(actual_relevances, predicted_relevances, k=None):
9
+ order = np.argsort(-predicted_relevances)
10
+ actual_relevances = actual_relevances[order]
11
+ if k is not None:
12
+ actual_relevances = actual_relevances[:k]
13
+ dcg = compute_dcg(actual_relevances)
14
+ idcg = compute_dcg(np.sort(actual_relevances)[::-1])
15
+ return dcg / idcg if idcg > 0 else 0