Update README.md

README.md

@@ -196,9 +196,9 @@ Place those in data/ and then:
 import numpy as np
 import json
 
-train_embeds = np.load("data/train_embeds.npy")      # shape (N, d)
+train_embeds = np.load("data/train_embeds.npy")
 with open("data/train_texts.json", "r") as f:
-    train_texts = json.load(f)                       # list[str]
+    train_texts = json.load(f)
 
 print(f"Loaded {train_embeds.shape[0]} embeddings of dim {train_embeds.shape[1]}")
 ```
@@ -220,13 +220,13 @@ batch = {
 with torch.no_grad():
     logits, metadata = model.proto_module(batch)
 
-attn_scores = metadata["attentions"][0]  # [num_labels, seq_len]
+attn_scores = metadata["attentions"][0]
 for label_id, scores in enumerate(attn_scores):
     topk = sorted(zip(batch["tokens"][0], scores.tolist()),
                   key=lambda x: -x[1])[:5]
     print(f"Label {label_id} top tokens:", topk)
 
-proto_vecs = model.proto_module.prototype_vectors.cpu().numpy()  # [num_labels, d]
+proto_vecs = model.proto_module.prototype_vectors.cpu().numpy()
 nn = NearestNeighbors(n_neighbors=1, metric="euclidean").fit(train_embeds)
 
 for label_id, u_c in enumerate(proto_vecs):