epsilon-compositionality

commit a8f598e73e82b30e8a8de93d39b70b63697a3752
parent 94346d57c6143ff2c5ca8584fa76825a68fe48f9
Author: Sondre Wold <sondrewold1@gmail.com>
Date:   Thu, 12 Jun 2025 15:51:44 +0200

Add selection of device, change output file name

Diffstat:
M
main.py
 | 
12
+++++++++---

1 file changed, 9 insertions(+), 3 deletions(-)
diff --git a/main.py b/main.py
@@ -12,6 +12,8 @@ import tqdm
 import torch
 import transformers
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 
 class Point(NamedTuple):
     word_similarity: float
@@ -22,10 +24,11 @@ class Point(NamedTuple):
 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.dataset_name = dataset["name"]
         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.model = transformers.AutoModel.from_pretrained(model).to(device)
         self.tokenizer = transformers.AutoTokenizer.from_pretrained(model)
 
     def join_tokens(self: Experiment, tokens: list[str]) -> str:
@@ -56,7 +59,7 @@ class Experiment:
             })
 
     def embed(self: Experiment, text: str) -> torch.Tensor:
-        return self.model(self.tokenizer(text, return_tensors="pt")["input_ids"]).pooler_output.cuda()
+        return self.model(self.tokenizer(text, return_tensors="pt")["input_ids"].to(device)).pooler_output
 
     def similarity(self: Experiment, lhs: str, rhs: str) -> float:
         # TODO cpu is needed?
@@ -80,10 +83,13 @@ class Experiment:
                         yield Point(word_similarity=self.similarity(first, second),
                                     contextualized_similarity=self.similarity(cfirst, csecond),
                                     weight=first_count/len(counter))
+                        break
+                    break
+                break
 
     def __call__(self: Experiment):
         points: list[Point] = list(self.all_similarities())
-        with open("hey.pkl", "wb") as file:
+        with open(f"points_output_{self.dataset_name}.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],