commit 94346d57c6143ff2c5ca8584fa76825a68fe48f9
Author: Sondre Wold <sondrewold1@gmail.com>
Date: Thu, 12 Jun 2025 15:43:58 +0200
add main file from Étienne
Diffstat:
| A | main.py | | | 107 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
1 file changed, 107 insertions(+), 0 deletions(-)
diff --git a/main.py b/main.py
@@ -0,0 +1,107 @@
+from __future__ import annotations
+from typing import Iterable, NamedTuple
+import collections
+import math
+import pickle
+import re
+
+import datasets
+import numpy
+import matplotlib.pyplot
+import tqdm
+import torch
+import transformers
+
+
+class Point(NamedTuple):
+ word_similarity: float
+ contextualized_similarity: float
+ weight: float
+
+
+class Experiment:
+ def __init__(self: Experiment, model: str, dataset: dict[str, str], ngram: int, grams_are_tokens: bool, minimal_occurences: int) -> None:
+ self.ngram = ngram
+ self.grams_are_tokens = grams_are_tokens
+ self.minimal_occurences = minimal_occurences
+ self.dataset = datasets.load_dataset(**dataset)
+ self.model = transformers.AutoModel.from_pretrained(model).cuda()
+ self.tokenizer = transformers.AutoTokenizer.from_pretrained(model)
+
+ def join_tokens(self: Experiment, tokens: list[str]) -> str:
+ if self.grams_are_tokens:
+ raise NotImplementedError("TODO tokenizer")
+ else:
+ return " ".join(tokens)
+
+ def loop_strings_to_ngrams(self: Experiment, input: list[str]) -> dict[str, collections.Counter]:
+ counters = collections.defaultdict(collections.Counter)
+ for i in range(len(input) - self.ngram - 1):
+ context: str = self.join_tokens(input[i+1:i + 1 + self.ngram])
+ if self.filter(input[i]) and self.filter(context):
+ counters[context][input[i]] += 1
+ return counters
+
+ def get_ngrams(self: Experiment) -> dict[str, dict[str, float]]:
+ counters = collections.defaultdict(collections.Counter)
+ for sample in tqdm.tqdm(self.dataset):
+ if self.grams_are_tokens:
+ raise NotImplementedError("TODO tokenizer")
+ else:
+ for key, value in self.loop_strings_to_ngrams(sample["text"].split()).items():
+ counters[key] += value
+ return collections.defaultdict(collections.Counter, {
+ key: value for key, value in counters.items()
+ if sum(value.values()) >= self.minimal_occurences and len(value) >= 2
+ })
+
+ def embed(self: Experiment, text: str) -> torch.Tensor:
+ return self.model(self.tokenizer(text, return_tensors="pt")["input_ids"]).pooler_output.cuda()
+
+ def similarity(self: Experiment, lhs: str, rhs: str) -> float:
+ # TODO cpu is needed?
+ return torch.nn.functional.cosine_similarity(self.embed(lhs), self.embed(rhs))[0].cpu().item()
+
+ def filter(self: Experiment, text: str) -> bool:
+ return bool(re.fullmatch(r"[a-zA-Z ]+", text))
+
+ def all_similarities(self: Experiment) -> Iterable[Point]:
+ ngrams = self.get_ngrams()
+ print(f"Number of {self.ngram}-grams after filtering: {len(ngrams)}")
+
+ with tqdm.tqdm(ngrams.items()) as ctqdm:
+ for context, counter in ctqdm:
+ ctqdm.set_description(context)
+ total_count: int = sum(counter.values())
+ for first, first_count in tqdm.tqdm(counter.items()):
+ cfirst = self.join_tokens([first, context])
+ for second, second_count in counter.items():
+ csecond = self.join_tokens([second, context])
+ yield Point(word_similarity=self.similarity(first, second),
+ contextualized_similarity=self.similarity(cfirst, csecond),
+ weight=first_count/len(counter))
+
+ def __call__(self: Experiment):
+ points: list[Point] = list(self.all_similarities())
+ with open("hey.pkl", "wb") as file:
+ pickle.dump(points, file)
+ data, xedges, yedges = numpy.histogram2d([point.word_similarity for point in points],
+ [point.contextualized_similarity for point in points],
+ weights=[point.weight for point in points],
+ bins=100)
+ extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
+ matplotlib.pyplot.clf()
+ matplotlib.pyplot.imshow(data.T, extent=extent, origin="lower")
+ matplotlib.pyplot.savefig("figure.svg", format="svg")
+ matplotlib.pyplot.show()
+
+
+if __name__ == "__main__":
+ Experiment(
+ model="bert-base-uncased",
+ dataset={"path":"Salesforce/wikitext", "name":"wikitext-2-v1", "split":"train"},
+ #dataset={"path":"HuggingFaceTB/smollm-corpus", "name": "fineweb-edu-dedup", "split":"train"},
+ ngram=3,
+ grams_are_tokens=False,
+ minimal_occurences=20,
+ )()
