File size: 3,532 Bytes
7e73a05 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import json
import numpy as np
import pickle
from tqdm import tqdm
import torch
import torch.nn as nn
def load_tokenizer():
"""Load the CBOW tokenizer mappings."""
print("Loading tokenizer...")
with open('tkn_words_to_ids.pkl', 'rb') as f:
words_to_ids = pickle.load(f)
with open('tkn_ids_to_words.pkl', 'rb') as f:
ids_to_words = pickle.load(f)
return words_to_ids, ids_to_words
def load_tokenized_triples():
"""Load the tokenized triples."""
print("Loading tokenized triples...")
with open('tokenized_triples.json', 'r') as f:
data = json.load(f)
return data
def create_embedding_layer(vocab_size, embedding_dim=128):
"""Create a simple embedding layer."""
embedding = nn.Embedding(vocab_size, embedding_dim)
# Initialize with random weights
nn.init.xavier_uniform_(embedding.weight)
return embedding
def average_pool(tokens, embedding_layer):
"""Create average pooled vector for a list of tokens."""
# Convert tokens to tensor
tokens_tensor = torch.tensor(tokens, dtype=torch.long)
# Get embeddings
embeddings = embedding_layer(tokens_tensor)
# Average the embeddings and detach before converting to numpy
return torch.mean(embeddings, dim=0).detach().numpy()
def process_triples(data, embedding_layer):
"""Process triples and create average pooled vectors."""
processed_data = {
'train': [],
'validation': [],
'test': []
}
for split in ['train', 'validation', 'test']:
print(f"\nProcessing {split} split...")
for triple in tqdm(data[split]):
# Get average pooled vectors
query_vector = average_pool(triple['query_tokens'], embedding_layer)
pos_doc_vector = average_pool(triple['positive_document_tokens'], embedding_layer)
neg_doc_vector = average_pool(triple['negative_document_tokens'], embedding_layer)
processed_data[split].append({
'query_vector': query_vector.tolist(),
'positive_document_vector': pos_doc_vector.tolist(),
'negative_document_vector': neg_doc_vector.tolist(),
'query': triple['query'], # Keep original text for reference
'positive_document': triple['positive_document'],
'negative_document': triple['negative_document']
})
return processed_data
def main():
# Load data
words_to_ids, ids_to_words = load_tokenizer()
data = load_tokenized_triples()
# Create embedding layer
vocab_size = len(words_to_ids)
embedding_layer = create_embedding_layer(vocab_size)
# Process triples
processed_data = process_triples(data, embedding_layer)
# Save processed data
print("\nSaving processed data...")
with open('triple_embeddings.json', 'w') as f:
json.dump(processed_data, f)
# Print statistics
for split in ['train', 'validation', 'test']:
print(f"\n{split.upper()} split:")
print(f"Number of processed triples: {len(processed_data[split])}")
if processed_data[split]:
sample = processed_data[split][0]
print("\nSample vector shapes:")
print("Query vector shape:", len(sample['query_vector']))
print("Positive doc vector shape:", len(sample['positive_document_vector']))
print("Negative doc vector shape:", len(sample['negative_document_vector']))
if __name__ == "__main__":
main() |