refactor: Modularize code by creating functions for model loading and processing

app.py

@@ -6,117 +6,68 @@ from pprint import pprint
 
 #%%
 
-model_name="mistralai/Mistral-7B-v0.1"
-
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-# model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Move the model to GPU if available
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-#%%
-
-input_text = "I just drive to the store to but eggs, but they had some."
-input_text = "He asked me to prostate myself before the king, but I rifused."
-input_text = "He asked me to prostrate myself before the king, but I rifused."
-
-#%%
-inputs = tokenizer(input_text, return_tensors="pt").to(device)
-input_ids = inputs["input_ids"]
-labels = input_ids
-
-#%%
-with torch.no_grad():
-    outputs = model(**inputs, labels=labels)
-
-#%%
-
-# Get logits and shift them
-logits = outputs.logits[0, :-1, :]
-
-# Calculate log probabilities
-log_probs = torch.log_softmax(logits, dim=-1)
-
-# Get the log probability of each token in the sequence
-token_log_probs = log_probs[range(log_probs.shape[0]), input_ids[0][1:]]
-
-# Decode tokens and pair them with their log probabilities
-tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
-result = list(zip(tokens[1:], token_log_probs.tolist()))
-
-#%%
-for token, logprob in result:
-    print(f"Token: {token}, Log Probability: {logprob:.4f}")
-
-# %%
-words = []
-current_word = []
-current_log_probs = []
-
-for token, logprob in result:
-    if not token.startswith(chr(9601)) and token.isalpha():
-        current_word.append(token)
-        current_log_probs.append(logprob)
-    else:
-        if current_word:
-            words.append(("".join(current_word), sum(current_log_probs)))
-        current_word = [token]
-        current_log_probs = [logprob]
-
-if current_word:
-    words.append(("".join(current_word), sum(current_log_probs)))
-
-for word, avg_logprob in words:
-    print(f"Word: {word}, Log Probability: {avg_logprob:.4f}")
-
-# %%
-
-
-words = split_into_words(tokens[1:], token_log_probs)
-
-# Define a threshold for low probability words
-log_prob_threshold = -5.0
-
-# Filter words with log probability below the threshold
-low_prob_words = [word for word in words if word.logprob < log_prob_threshold]
-
-
-#%%
-def generate_replacements(model, tokenizer, prefix, num_samples=5):
+def load_model_and_tokenizer(model_name):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    return model, tokenizer, device
+
+def process_input_text(input_text, tokenizer, device):
+    inputs = tokenizer(input_text, return_tensors="pt").to(device)
+    input_ids = inputs["input_ids"]
+    return inputs, input_ids
+
+def calculate_log_probabilities(model, inputs, input_ids):
+    with torch.no_grad():
+        outputs = model(**inputs, labels=input_ids)
+    logits = outputs.logits[0, :-1, :]
+    log_probs = torch.log_softmax(logits, dim=-1)
+    token_log_probs = log_probs[range(log_probs.shape[0]), input_ids[0][1:]]
+    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
+    return list(zip(tokens[1:], token_log_probs.tolist()))
+
+def generate_replacements(model, tokenizer, prefix, device, num_samples=5):
     input_context = tokenizer(prefix, return_tensors="pt").to(device)
     input_ids = input_context["input_ids"]
-
     new_words = []
     for _ in range(num_samples):
         with torch.no_grad():
             outputs = model.generate(
                 input_ids=input_ids,
-                max_length=input_ids.shape[-1] + 5,  # generate a few tokens beyond the prefix
+                max_length=input_ids.shape[-1] + 5,
                 num_return_sequences=1,
                 temperature=1.0,
-                top_k=50,     # use top-k sampling
-                top_p=0.95,   # use nucleus sampling
+                top_k=50,
+                top_p=0.95,
                 do_sample=True
             )
-
-        generated_ids = outputs[0][input_ids.shape[-1]:]  # extract the newly generated part
+        generated_ids = outputs[0][input_ids.shape[-1]:]
         new_word = tokenizer.decode(generated_ids, skip_special_tokens=True).split()[0]
         new_words.append(new_word)
-
     return new_words
 
-# Generate new words for low probability words
-for word in low_prob_words:
-    prefix_index = word.first_token_index
-    prefix_tokens = tokens[:prefix_index + 1]  # include the word itself
-    prefix = tokenizer.convert_tokens_to_string(prefix_tokens)
+def main():
+    model_name = "mistralai/Mistral-7B-v0.1"
+    model, tokenizer, device = load_model_and_tokenizer(model_name)
+
+    input_text = "He asked me to prostrate myself before the king, but I rifused."
+    inputs, input_ids = process_input_text(input_text, tokenizer, device)
+
+    result = calculate_log_probabilities(model, inputs, input_ids)
 
-    replacements = generate_replacements(model, tokenizer, prefix)
+    words = split_into_words([token for token, _ in result], [logprob for _, logprob in result])
+    log_prob_threshold = -5.0
+    low_prob_words = [word for word in words if word.logprob < log_prob_threshold]
 
-    print(f"Original word: {word.text}, Log Probability: {word.logprob:.4f}")
-    print(f"Proposed replacements: {replacements}")
-    print()
+    for word in low_prob_words:
+        prefix_index = word.first_token_index
+        prefix_tokens = [token for token, _ in result][:prefix_index + 1]
+        prefix = tokenizer.convert_tokens_to_string(prefix_tokens)
+        replacements = generate_replacements(model, tokenizer, prefix, device)
+        print(f"Original word: {word.text}, Log Probability: {word.logprob:.4f}")
+        print(f"Proposed replacements: {replacements}")
+        print()
 
-# %%
+if __name__ == "__main__":
+    main()