Update app.py

app.py

@@ -140,8 +140,8 @@ def run_interpretation(raw_original_prompt, raw_interpretation_prompt, max_new_t
     # try identifying important layers
     vectors_to_compare = interpreted_vectors # torch.tensor(global_state.sentence_transformer.encode(generation_texts))
     diff_score1 = F.normalize(vectors_to_compare, dim=-1).diff(dim=0).norm(dim=-1).cpu()
-    
-    bags_of_words = [set(tokenizer.tokenize(text)) for text in generation_texts]
+    tokenized_generations = [tokenizer.tokenize(text) for text in generation_texts]
+    bags_of_words = [set(tokens) | set([(tokens[i], tokens[i+1]) for i in range(len(tokens)-1)]) for tokens in tokenized_generations]
     diff_score2 = torch.tensor([
                                 -len(bags_of_words[i+1] & bags_of_words[i]) / np.sqrt(len(bags_of_words[i+1]) * len(bags_of_words[i]))
                                 for i in range(len(bags_of_words)-1)