init model
Browse files- LICENSE +21 -0
- data_processing.py +55 -0
- inference.py +48 -0
- label_mappings.json +1 -0
- multioutput_regressor.pth +3 -0
- requirements.txt +7 -0
- train.py +301 -0
- 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
|