133 lines
4.7 KiB
Python
133 lines
4.7 KiB
Python
from typing import Tuple
|
|
import re
|
|
from html import unescape
|
|
import torch
|
|
from torch.utils.data import Dataset
|
|
from datasets import load_dataset # Hugging Face (not PyTorch)
|
|
import scipy
|
|
import numpy as np
|
|
# Monkey-patch gensim as it relies on scipy's deprecated triu method
|
|
# /gensim/matutils.py has this line which causes ImportError:
|
|
# from scipy.linalg import get_blas_funcs, triu
|
|
# TODO: Monitor gensim and fix this if upgrade is made
|
|
scipy.linalg.triu = np.triu
|
|
import gensim.downloader as api
|
|
|
|
class HfGloveDataset(Dataset):
|
|
"""
|
|
Hugging Face Dataset wrapper with word embedding support.
|
|
|
|
pragma: skip_doc
|
|
"""
|
|
def __init__(self,
|
|
dataset_name: str,
|
|
split: str,
|
|
glove_dim: int = 100,
|
|
max_seq_len: int = 256,
|
|
sample_field_name: str = "text",
|
|
label_field_name: str = "label") -> None:
|
|
"""
|
|
Initializes the dataset by loading Hugging Face dataset and GloVe vectors.
|
|
|
|
Args:
|
|
dataset_name (str): The name of the Hugging Face dataset.
|
|
split (str): The dataset split (e.g., 'train').
|
|
glove_dim (int): Dimensionality of GloVe embeddings.
|
|
max_seq_len (int): Maximum sequence length.
|
|
sample_field_name (str): Field name for text samples.
|
|
label_field_name (str): Field name for labels.
|
|
"""
|
|
self.dataset = load_dataset(dataset_name, split=split)
|
|
self.glove_dim = glove_dim
|
|
self.max_seq_len = max_seq_len
|
|
self.sample_field_name = sample_field_name
|
|
self.label_field_name = label_field_name
|
|
|
|
# Load GloVe from gensim
|
|
glove_name = f"glove-wiki-gigaword-{glove_dim}"
|
|
self.glove = api.load(glove_name)
|
|
self.unk_vector = torch.zeros(glove_dim)
|
|
|
|
def tokenize(self, text: str) -> list[str]:
|
|
"""
|
|
Tokenizes raw text into a list of lowercase word tokens.
|
|
It removes HTML tags, decodes HTML entities, strips punctuation
|
|
(while keeping letters and numbers), and splits on whitespace.
|
|
|
|
Args:
|
|
text (str): Raw text input (may include HTML and punctuation).
|
|
|
|
Returns:
|
|
list[str]: List of clean, lowercase word tokens.
|
|
"""
|
|
# Decode HTML entities like "&", "<", etc. to their actual characters
|
|
text = unescape(text)
|
|
|
|
# Remove HTML tags (like <br />, <p>, etc.)
|
|
# The regex "<.*?>" uses .*? (non-greedy match) to ensure it matches one tag at a time,
|
|
# instead of greedily removing everything between the first '<' and last '>'
|
|
text = re.sub(r"<.*?>", " ", text)
|
|
|
|
# Remove all characters except letters, digits, and whitespace
|
|
# This strips punctuation like "!", ".", "(", but keeps words and numbers
|
|
text = re.sub(r"[^a-zA-Z0-9\s]", " ", text)
|
|
|
|
# Replace consecutive white space to a single white space
|
|
text = re.sub(r" +", " ", text)
|
|
|
|
# Convert to lowercase and split by whitespace to get token list
|
|
return text.lower().split()
|
|
|
|
|
|
def vectorize(self, tokens: list[str]) -> torch.Tensor:
|
|
"""
|
|
Converts tokens to a fixed-length tensor of word vectors.
|
|
|
|
Args:
|
|
tokens (list[str]): A list of word tokens.
|
|
|
|
Returns:
|
|
torch.Tensor: A tensor of shape (max_seq_len, glove_dim).
|
|
"""
|
|
vectors = []
|
|
for token in tokens[:self.max_seq_len]:
|
|
if token in self.glove:
|
|
vec = torch.tensor(self.glove[token], dtype=torch.float32)
|
|
else:
|
|
vec = self.unk_vector
|
|
vectors.append(vec)
|
|
|
|
if len(vectors) < self.max_seq_len:
|
|
pad = [torch.zeros(self.glove_dim) for _ in range(self.max_seq_len - len(vectors))]
|
|
vectors.extend(pad)
|
|
|
|
return torch.stack(vectors)
|
|
|
|
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Retrieves the (embedded sequence, label) pair at the given index.
|
|
|
|
Args:
|
|
idx (int): Index of the item.
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor, torch.Tensor]: Embedding tensor and label tensor.
|
|
"""
|
|
# Get a text sample from database and convert to a word embedding
|
|
sample = self.dataset[idx]
|
|
tokens = self.tokenize(sample[self.sample_field_name])
|
|
emb_tensor = self.vectorize(tokens)
|
|
|
|
# Get a corresponding label
|
|
label = torch.tensor(sample[self.label_field_name], dtype=torch.long)
|
|
return emb_tensor, label
|
|
|
|
def __len__(self) -> int:
|
|
"""
|
|
Returns the number of samples in the dataset.
|
|
|
|
Returns:
|
|
int: Dataset length.
|
|
"""
|
|
return len(self.dataset)
|