From e1b85ebb5d588499f74bdac70cc4724840ef1b2d Mon Sep 17 00:00:00 2001
From: Pavan Kumar <66913595+ascender1729@users.noreply.github.com>
Date: Fri, 13 Jun 2025 21:51:16 +0530
Subject: [PATCH] docs: add logits processor guide

---
 docs/mkdocs.yml                       |  1 +
 docs/swarms/tools/logits_processor.md | 55 +++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)
 create mode 100644 docs/swarms/tools/logits_processor.md

diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml
index d4e9725c..6bf5f8d9 100644
--- a/docs/mkdocs.yml
+++ b/docs/mkdocs.yml
@@ -267,6 +267,7 @@ nav:
       - Overview: "swarms_tools/overview.md"
       - BaseTool Reference: "swarms/tools/base_tool.md"
       - MCP Client Utils: "swarms/tools/mcp_client_call.md"
+      - Logits Processor: "swarms/tools/logits_processor.md"
 
       - Vertical Tools:
         - Finance: "swarms_tools/finance.md"
diff --git a/docs/swarms/tools/logits_processor.md b/docs/swarms/tools/logits_processor.md
new file mode 100644
index 00000000..04ca0689
--- /dev/null
+++ b/docs/swarms/tools/logits_processor.md
@@ -0,0 +1,55 @@
+# Logits Processor
+
+The `logits_processor` module offers utility classes to control token generation
+when using Hugging Face Transformers models. These processors can limit the
+model to specific tokens or halt generation once a condition is met. They are
+particularly helpful when a model must output well‑formed numbers or when you
+need to stop after producing certain characters.
+
+## Installation
+
+The classes rely on `torch` and `transformers`. They are automatically installed
+when you import the module, but you can install them manually:
+
+```bash
+pip install torch transformers
+```
+
+## Classes
+
+| Class | Description |
+|-------|-------------|
+| `StringStoppingCriteria` | Stops generation after the prompt when a `"` character is produced. |
+| `NumberStoppingCriteria` | Prevents invalid or overly precise numbers and stops when a complete number is detected. |
+| `OutputNumbersTokens` | A `LogitsWarper` that masks all tokens except digits and decimal points so only numbers can be generated. |
+
+## Basic Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from swarms.tools.logits_processor import (
+    OutputNumbersTokens,
+    NumberStoppingCriteria,
+)
+
+model = AutoModelForCausalLM.from_pretrained("gpt2")
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+
+prompt = "The result is "
+input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+
+logit_processor = OutputNumbersTokens(tokenizer, prompt)
+stoppers = [NumberStoppingCriteria(tokenizer, len(input_ids[0]))]
+
+output = model.generate(
+    input_ids,
+    logits_processor=[logit_processor],
+    stopping_criteria=stoppers,
+    max_new_tokens=6,
+)
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+The example above forces the model to emit a short numeric value after the
+prompt. You can mix and match the processors to adapt them to your use case.
+