|
|
|
import json |
|
import torch |
|
from transformers import DistilBertTokenizerFast, DistilBertModel |
|
import numpy as np |
|
|
|
def load_data(file_path): |
|
with open(file_path, 'r') as f: |
|
dataset = json.load(f) |
|
outdata = [ |
|
{ |
|
"did": e["user_id"], |
|
"description": e["description"], |
|
"label_weights": e["user_categories"] |
|
} |
|
for e in dataset |
|
if e["description"] and e["user_categories"] |
|
] |
|
return outdata |
|
|
|
def prepare_labels(outdata): |
|
all_labels = sorted({label for record in outdata for label in record['label_weights'].keys()}) |
|
label2id = {label: i for i, label in enumerate(all_labels)} |
|
id2label = {i: label for label, i in label2id.items()} |
|
|
|
y_matrix = np.zeros((len(outdata), len(all_labels)), dtype=float) |
|
for idx, record in enumerate(outdata): |
|
for label, weight in record['label_weights'].items(): |
|
y_matrix[idx, label2id[label]] = weight |
|
return y_matrix, label2id, id2label |
|
|
|
class EmbeddingGenerator: |
|
def __init__(self, model_name='distilbert-base-uncased', device=None): |
|
self.tokenizer = DistilBertTokenizerFast.from_pretrained(model_name) |
|
self.embedding_model = DistilBertModel.from_pretrained(model_name) |
|
self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
self.embedding_model.to(self.device) |
|
|
|
def generate_embeddings(self, descriptions, batch_size=1000): |
|
all_embeddings = [] |
|
descriptions = [desc for desc in descriptions] |
|
for i in range(0, len(descriptions), batch_size): |
|
batch_descriptions = descriptions[i:i + batch_size] |
|
inputs = self.tokenizer( |
|
batch_descriptions, |
|
padding=True, |
|
truncation=True, |
|
max_length=128, |
|
return_tensors="pt" |
|
).to(self.device) |
|
with torch.no_grad(): |
|
outputs = self.embedding_model(**inputs) |
|
batch_embeddings = outputs.last_hidden_state[:, 0, :].cpu().numpy() |
|
all_embeddings.append(batch_embeddings) |
|
return np.vstack(all_embeddings) |
|
|
