Update README.md

README.md

@@ -58,6 +58,9 @@ def wrap_prompt(prompt: str, tokenizer) -> str:
         return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
 def detect_gene_token_space(tokenizer):
+  """
+  Used to detect the boundary of special gene tokens in the tokenizer vocabulary for subsequent masking non-gene tokens.
+  """
     vocab = tokenizer.get_vocab()
     gene_ids = sorted(idx for tok, idx in vocab.items() if tok.startswith(GENE_PREFIX))
     if not gene_ids:
@@ -76,12 +79,21 @@ with torch.no_grad():
     inputs = tokenizer([text], return_tensors="pt")
     inputs = {k: v.to(next(model.parameters()).device) for k, v in inputs.items()}
     outputs = model(**inputs)
+
     # Take logits for the last token only
     logits = outputs.logits[0, -1, :].to("cpu", dtype=torch.float32)
+
     # Mask non-gene tokens
     gene_token_tensor = gene_token_tensor.to(dtype=torch.long, device=logits.device)
     gene_logits = logits[gene_token_tensor]
+
+    # Sort by logits
     order = torch.argsort(gene_logits, descending=True)
+
+    # Ranked Gene List with Scores (logits)
     sorted_gene_ids = gene_token_tensor[order]
     sorted_gene_tok = tokenizer.convert_ids_to_tokens(sorted_gene_ids.tolist())
-```
+```
+
+## Special Gene Token Example
+In `<|PhenoGemini-Special-Token-Entrez-ID-64087|>`, the numeric value after the final hyphen represents the Entrez Gene ID.