Snapshot

app.py

@@ -1,14 +1,52 @@
 #%%
 import time
-from text_processing import split_into_words, Word
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast, BatchEncoding
+from transformers.generation.utils import GenerateOutput
 from typing import cast
+from dataclasses import dataclass
 
 type Tokenizer = PreTrainedTokenizer | PreTrainedTokenizerFast
 
+@dataclass
+class Word:
+    tokens: list[int]
+    text: str
+    logprob: float
+    context: list[int]
+
+def split_into_words(token_probs: list[tuple[int, float]], tokenizer: Tokenizer) -> list[Word]:
+    words: list[Word] = []
+    current_word: list[int] = []
+    current_log_probs: list[float] = []
+    current_word_first_token_index: int = 0
+    all_tokens: list[int] = [token_id for token_id, _ in token_probs]
+
+    def append_current_word():
+        if current_word:
+            words.append(Word(current_word,
+                              tokenizer.decode(current_word),
+                              sum(current_log_probs),
+                              all_tokens[:current_word_first_token_index]))
+
+    for i, (token_id, logprob) in enumerate(token_probs):
+        token: str = tokenizer.convert_ids_to_tokens([token_id])[0]
+        if not token.startswith(chr(9601)) and token.isalpha():
+            current_word.append(token_id)
+            current_log_probs.append(logprob)
+        else:
+            append_current_word()
+            current_word = [token_id]
+            current_log_probs = [logprob]
+            current_word_first_token_index = i
+
+    append_current_word()
+
+    return words
+
 def load_model_and_tokenizer(model_name: str, device: torch.device) -> tuple[PreTrainedModel, Tokenizer]:
-    tokenizer: Tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer: Tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
+    tokenizer.pad_token = tokenizer.eos_token
     model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(model_name)
     model.to(device)
     return model, tokenizer
@@ -32,10 +70,16 @@ def calculate_log_probabilities(model: PreTrainedModel, tokenizer: Tokenizer, in
     tokens: torch.Tensor = input_ids[0][1:]
     return list(zip(tokens.tolist(), token_log_probs.tolist()))
 
+def prepare_inputs(contexts: list[list[int]], tokenizer: Tokenizer, device: torch.device) -> tuple[torch.FloatTensor, torch.FloatTensor]:
+    texts = [tokenizer.decode(context, skip_special_tokens=True) for context in contexts]
+    inputs = tokenizer(texts, return_tensors="pt", padding=True).to(device)
+    input_ids = cast(torch.FloatTensor, inputs["input_ids"])
+    attention_mask = cast(torch.FloatTensor, inputs["attention_mask"])
+    return input_ids, attention_mask
 
-def generate_replacements(model: PreTrainedModel, tokenizer: Tokenizer, contexts: list[list[int]], device: torch.device, num_samples: int = 5) -> list[list[str]]:
-    input_ids = torch.tensor(contexts).to(device)
-    attention_mask = torch.ones_like(input_ids)
+def generate_replacements(model: PreTrainedModel, tokenizer: Tokenizer, contexts: list[list[int]],
+                          device: torch.device, num_samples: int = 5) -> tuple[GenerateOutput | torch.LongTensor, list[list[str]]]:
+    input_ids, attention_mask = prepare_inputs(contexts, tokenizer, device)
     with torch.no_grad():
         outputs = model.generate(
             input_ids=input_ids,
@@ -52,10 +96,10 @@ def generate_replacements(model: PreTrainedModel, tokenizer: Tokenizer, contexts
         replacements = []
         for j in range(num_samples):
             generated_ids = outputs[i * num_samples + j][input_ids.shape[-1]:]
-            new_word = tokenizer.decode(generated_ids, skip_special_tokens=True).split()[0]
+            new_word = tokenizer.decode(generated_ids, skip_special_tokens=False).split()[0]
             replacements.append(new_word)
         all_new_words.append(replacements)
-    return all_new_words
+    return outputs, all_new_words
 
 #%%
 
@@ -75,28 +119,24 @@ token_probs: list[tuple[int, float]] = calculate_log_probabilities(model, tokeni
 
 #%%
 
-import importlib
-import text_processing
-
-importlib.reload(text_processing)
-from text_processing import split_into_words, Word
-
 words = split_into_words(token_probs, tokenizer)
 log_prob_threshold = -5.0
 low_prob_words = [word for word in words if word.logprob < log_prob_threshold]
 
+#%%
+contexts = [word.context for word in low_prob_words]
+
 #%%
 
 start_time = time.time()
+replacements_batch = generate_replacements(model, tokenizer, contexts, device, num_samples=5)
+end_time = time.time()
+print(f"Total time taken for replacements: {end_time - start_time:.4f} seconds")
 
-contexts = [word.context for word in low_prob_words]
-replacements_batch = generate_replacements(model, tokenizer, contexts, device)
+#%%
 
 for word, replacements in zip(low_prob_words, replacements_batch):
     print(f"Original word: {word.text}, Log Probability: {word.logprob:.4f}")
     print(f"Proposed replacements: {replacements}")
 
-end_time = time.time()
-print(f"Total time taken for replacements: {end_time - start_time:.4f} seconds")
-
 # %%

text_processing.py

@@ -1,40 +0,0 @@
-from dataclasses import dataclass
-from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
-
-type Tokenizer = PreTrainedTokenizer | PreTrainedTokenizerFast
-
-@dataclass
-class Word:
-    tokens: list[int]
-    text: str
-    logprob: float
-    context: list[int]
-
-def split_into_words(token_probs: list[tuple[int, float]], tokenizer: Tokenizer) -> list[Word]:
-    words: list[Word] = []
-    current_word: list[int] = []
-    current_log_probs: list[float] = []
-    current_word_first_token_index: int = 0
-    all_tokens: list[int] = [token_id for token_id, _ in token_probs]
-
-    def append_current_word():
-        if current_word:
-            words.append(Word(current_word,
-                              tokenizer.decode(current_word),
-                              sum(current_log_probs),
-                              all_tokens[:current_word_first_token_index]))
-
-    for i, (token_id, logprob) in enumerate(token_probs):
-        token: str = tokenizer.convert_ids_to_tokens([token_id])[0]
-        if not token.startswith(chr(9601)) and token.isalpha():
-            current_word.append(token_id)
-            current_log_probs.append(logprob)
-        else:
-            append_current_word()
-            current_word = [token_id]
-            current_log_probs = [logprob]
-            current_word_first_token_index = i
-
-    append_current_word()
-
-    return words