red-candle 1.0.2 → 1.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 7405c9911d6088106dd7a19e96312f12b86e9a80087c3d7745cd3911e263890a
-  data.tar.gz: a88b75152708e72e019aba9acfeabd899f1ae1d1c567562ded6d2c6aa8eae8d0
+  metadata.gz: 5a2b4fac15c2c261d8fea90d34973605209043e2b2a222f82414bf94bc5c47e8
+  data.tar.gz: daff4398f34170e20744ab2ee8b1abb5b977ffec15a129d678efced7c9649495
 SHA512:
-  metadata.gz: ce1cc52dc1223968f3398ab0972283a6309a80d306c14193f23336cd36ed55c8fa5eaaaf05d756f76c88e442abe19b0d82d2742a49199930ef6effcffd6d4482
-  data.tar.gz: 8c30f3c0c096f8186b219a9a5d0fe92928621126f545eaabae26ddedc843515b7da8c45890a1f24a7d519c0c68fcd55b0553ea73f977644258ff30c5e5ccd2f1
+  metadata.gz: 4391a7fb4072d9ac174bcecdb366c975c5e487f43fcc4b0db75533aa44c94822d38103a633eba0b098ac26cf31313e9dd7da77255ed83692584e6936cca86271
+  data.tar.gz: da7a15a86fea349069079537b8c6f6696079842d205168aa908644a59528d3ecf5922548d8b3d153494f0beb3fe6221b30110368f1f1289f0ac1ff4856f6b243

data/Cargo.lock

@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 3
+version = 4
 
 [[package]]
 name = "accelerate-src"
@@ -121,6 +121,26 @@ version = "0.22.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
 
+[[package]]
+name = "bincode"
+version = "2.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "36eaf5d7b090263e8150820482d5d93cd964a81e4019913c972f4edcc6edb740"
+dependencies = [
+ "bincode_derive",
+ "serde",
+ "unty",
+]
+
+[[package]]
+name = "bincode_derive"
+version = "2.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf95709a440f45e986983918d0e8a1f30a9b1df04918fc828670606804ac3c09"
+dependencies = [
+ "virtue",
+]
+
 [[package]]
 name = "bindgen"
 version = "0.69.5"
@@ -136,7 +156,7 @@ dependencies = [
  "proc-macro2",
  "quote",
  "regex",
- "rustc-hash",
+ "rustc-hash 1.1.0",
  "shlex",
  "syn",
 ]
@@ -255,13 +275,14 @@ dependencies = [
  "candle-nn",
  "candle-transformers",
  "half",
- "hf-hub",
+ "hf-hub 0.4.3",
  "magnus",
+ "outlines-core",
  "rand 0.8.5",
  "safetensors 0.3.3",
  "serde",
  "serde_json",
- "tokenizers",
+ "tokenizers 0.21.2",
  "tokio",
 ]
 
@@ -641,6 +662,15 @@ dependencies = [
  "dirs-sys 0.4.1",
 ]
 
+[[package]]
+name = "dirs"
+version = "5.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "44c45a9d03d6676652bcb5e724c7e988de1acad23a711b5217ab9cbecbec2225"
+dependencies = [
+ "dirs-sys 0.4.1",
+]
+
 [[package]]
 name = "dirs"
 version = "6.0.0"
@@ -1303,13 +1333,30 @@ version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c"
 
+[[package]]
+name = "hf-hub"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2b780635574b3d92f036890d8373433d6f9fc7abb320ee42a5c25897fc8ed732"
+dependencies = [
+ "dirs 5.0.1",
+ "indicatif",
+ "log",
+ "native-tls",
+ "rand 0.8.5",
+ "serde",
+ "serde_json",
+ "thiserror 1.0.69",
+ "ureq",
+]
+
 [[package]]
 name = "hf-hub"
 version = "0.4.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97"
 dependencies = [
- "dirs",
+ "dirs 6.0.0",
  "futures",
  "http",
  "indicatif",
@@ -1605,6 +1652,12 @@ dependencies = [
  "web-time",
 ]
 
+[[package]]
+name = "indoc"
+version = "2.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f4c7245a08504955605670dbf141fceab975f15ca21570696aebe9d2e71576bd"
+
 [[package]]
 name = "intel-mkl-src"
 version = "0.8.1"
@@ -1654,6 +1707,15 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "itertools"
+version = "0.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "itertools"
 version = "0.12.1"
@@ -1815,6 +1877,15 @@ dependencies = [
  "stable_deref_trait",
 ]
 
+[[package]]
+name = "memoffset"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
+dependencies = [
+ "autocfg",
+]
+
 [[package]]
 name = "metal"
 version = "0.27.0"
@@ -2186,6 +2257,25 @@ version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d"
 
+[[package]]
+name = "outlines-core"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4f0964d94d3e2322d2c0bbf80549affe085e2c6df08cf6c06e8c558988bcb11b"
+dependencies = [
+ "bincode",
+ "hf-hub 0.3.2",
+ "once_cell",
+ "regex",
+ "regex-automata",
+ "rustc-hash 2.1.1",
+ "serde",
+ "serde-pyobject",
+ "serde_json",
+ "thiserror 2.0.12",
+ "tokenizers 0.20.3",
+]
+
 [[package]]
 name = "paste"
 version = "1.0.15"
@@ -2306,6 +2396,69 @@ dependencies = [
  "version_check",
 ]
 
+[[package]]
+name = "pyo3"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884"
+dependencies = [
+ "cfg-if",
+ "indoc",
+ "libc",
+ "memoffset",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+ "unindent",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38"
+dependencies = [
+ "once_cell",
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.22.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "pyo3-build-config",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "quote"
 version = "1.0.40"
@@ -2418,6 +2571,17 @@ dependencies = [
  "rayon-core",
 ]
 
+[[package]]
+name = "rayon-cond"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "059f538b55efd2309c9794130bc149c6a553db90e9d99c2030785c82f0bd7df9"
+dependencies = [
+ "either",
+ "itertools 0.11.0",
+ "rayon",
+]
+
 [[package]]
 name = "rayon-cond"
 version = "0.4.0"
@@ -2604,6 +2768,12 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
 
+[[package]]
+name = "rustc-hash"
+version = "2.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d"
+
 [[package]]
 name = "rustix"
 version = "1.0.7"
@@ -2740,6 +2910,16 @@ dependencies = [
  "serde_derive",
 ]
 
+[[package]]
+name = "serde-pyobject"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ca4b0aad8b225845739a0030a0d5cc2ae949c56a86a7daf9226c7df7c2016d16"
+dependencies = [
+ "pyo3",
+ "serde",
+]
+
 [[package]]
 name = "serde_derive"
 version = "1.0.219"
@@ -2757,6 +2937,7 @@ version = "1.0.140"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "20068b6e96dc6c9bd23e01df8827e6c7e1f2fddd43c21810382803c136b99373"
 dependencies = [
+ "indexmap",
  "itoa",
  "memchr",
  "ryu",
@@ -2995,6 +3176,12 @@ dependencies = [
  "xattr",
 ]
 
+[[package]]
+name = "target-lexicon"
+version = "0.12.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
+
 [[package]]
 name = "tempfile"
 version = "3.20.0"
@@ -3058,6 +3245,39 @@ dependencies = [
  "zerovec",
 ]
 
+[[package]]
+name = "tokenizers"
+version = "0.20.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "67b67c92f6d705e2a1d106fb0b28c696f9074901a9c656ee5d9f5de204c39bf7"
+dependencies = [
+ "aho-corasick",
+ "derive_builder",
+ "esaxx-rs",
+ "getrandom 0.2.16",
+ "hf-hub 0.3.2",
+ "indicatif",
+ "itertools 0.12.1",
+ "lazy_static",
+ "log",
+ "macro_rules_attribute",
+ "monostate",
+ "onig",
+ "paste",
+ "rand 0.8.5",
+ "rayon",
+ "rayon-cond 0.3.0",
+ "regex",
+ "regex-syntax",
+ "serde",
+ "serde_json",
+ "spm_precompiled",
+ "thiserror 1.0.69",
+ "unicode-normalization-alignments",
+ "unicode-segmentation",
+ "unicode_categories",
+]
+
 [[package]]
 name = "tokenizers"
 version = "0.21.2"
@@ -3081,7 +3301,7 @@ dependencies = [
  "paste",
  "rand 0.9.1",
  "rayon",
- "rayon-cond",
+ "rayon-cond 0.4.0",
  "regex",
  "regex-syntax",
  "serde",
@@ -3365,12 +3585,24 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
 
+[[package]]
+name = "unindent"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
+
 [[package]]
 name = "untrusted"
 version = "0.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
 
+[[package]]
+name = "unty"
+version = "0.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d49784317cd0d1ee7ec5c716dd598ec5b4483ea832a2dced265471cc0f690ae"
+
 [[package]]
 name = "ureq"
 version = "2.12.1"
@@ -3431,6 +3663,12 @@ version = "0.9.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
 
+[[package]]
+name = "virtue"
+version = "0.0.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "051eb1abcf10076295e815102942cc58f9d5e3b4560e46e53c21e8ff6f3af7b1"
+
 [[package]]
 name = "walkdir"
 version = "2.5.0"

data/README.md

@@ -3,7 +3,7 @@
 [![build](https://github.com/assaydepot/red-candle/actions/workflows/build.yml/badge.svg)](https://github.com/assaydepot/red-candle/actions/workflows/build.yml)
 [![Gem Version](https://badge.fury.io/rb/red-candle.svg)](https://badge.fury.io/rb/red-candle)
 
-Run state-of-the-art **language models directly from Ruby**. No Python, no APIs, no external services - just Ruby with blazing-fast Rust under the hood. Hardware accelerated with **Metal (Mac)** and **CUDA (NVIDIA).**
+Run state-of-the-art **language models directly from Ruby**. No Python, no APIs, no external services - just Ruby with blazing-fast Rust under the hood. Hardware accelerated with **Metal (Mac)** and **CUDA (NVIDIA).** Red candle leverages the Rust ecosystem, notably [Candle](https://github.com/huggingface/candle) and [Magnus](https://github.com/matsadler/magnus), to provide a fast and efficient way to run LLMs in Ruby. See [Dependencies](#dependencies) for more.
 
 ## Install & Chat in 30 Seconds
 
@@ -126,6 +126,10 @@ Red-Candle now supports Large Language Models (LLMs) with GPU acceleration!
 - **Gemma**: Google's Gemma models (e.g., `google/gemma-2b`, `google/gemma-7b`, `google/gemma-2b-it`)
 - **Llama**: Llama 2 and Llama 3 models (e.g., `TinyLlama/TinyLlama-1.1B-Chat-v1.0`, `meta-llama/Llama-2-7b-hf`, `NousResearch/Llama-2-7b-hf`)
 - **Mistral**: All Mistral models (e.g., `mistralai/Mistral-7B-Instruct-v0.1`)
+- **Qwen**: Qwen 2 and 2.5 models (e.g., `Qwen/Qwen2-1.5B`, `Qwen/Qwen2.5-7B-Instruct`)
+- **Phi**: Microsoft's Phi-2, Phi-3, Phi-3.5, and Phi-4 models (e.g., `microsoft/phi-2`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/phi-4`)
+  - ⚠️ ⚠️ ⚠️ Note: Phi-3 and Phi-4 GGUF models have a known issue with KV cache persistence between generations. The `reset_cache` parameter doesn't work for GGUF models. Recreate the model instance for each generation.
+  - `candle` pull request about phi-3 gguf models: https://github.com/huggingface/candle/pull/2937
 
 ### Quantized Model Support (GGUF)
 
@@ -244,6 +248,36 @@ This is particularly useful for:
 - Troubleshooting generation problems
 - Analyzing model behavior
 
+## Structured Generation
+
+Red Candle supports structured generation to constrain LLM outputs to follow specific patterns like JSON schemas or regular expressions:
+
+```ruby
+# Define a JSON schema
+schema = {
+  type: "object",
+  properties: {
+    answer: { type: "string", enum: ["yes", "no"] },
+    confidence: { type: "number", minimum: 0, maximum: 1 }
+  },
+  required: ["answer"]
+}
+
+# Generate and parse in one step
+result = llm.generate_structured("Is Ruby easy to learn?", schema: schema)
+puts result["answer"]      # "yes"
+puts result["confidence"]  # 0.9
+
+# Or use regex patterns for non-JSON outputs
+phone_constraint = llm.constraint_from_regex('\d{3}-\d{3}-\d{4}')
+config = Candle::GenerationConfig.balanced(constraint: phone_constraint)
+phone = llm.generate("Generate a phone number:", config: config)
+```
+
+See [STRUCTURED_GENERATION.md](docs/STRUCTURED_GENERATION.md) for detailed documentation.
+
+**Note on Reliability**: Structured generation constrains the model's output tokens, but success rates vary by model size and schema complexity. Smaller models (< 7B parameters) may occasionally produce incomplete or invalid JSON, especially with complex schemas. Consider implementing retry logic or fallback strategies in production applications. Larger models generally perform much better with structured generation.
+
 ## ⚠️ Model Format Requirements
 
 ### EmbeddingModels and Rerankers: Safetensors Only
@@ -861,7 +895,7 @@ Pull requests are welcome.
 4. `git push --follow-tags`
 5. `gem push pkg/red-candle-VERSION_NUMBER.gem`
 
-## See Also
+## Dependencies
 
 - [Candle](https://github.com/huggingface/candle)
 - [Magnus](https://github.com/matsadler/magnus)

data/Rakefile

@@ -8,7 +8,14 @@ task default: :test
 Rake::TestTask.new do |t|
   t.deps << :compile
   t.libs << "test"
-  t.test_files = FileList["test/**/*_test.rb"].exclude("test/benchmarks/**/*_test.rb")
+  t.test_files = FileList["test/**/*_test.rb"]
+    .exclude("test/benchmarks/**/*_test.rb")
+    .exclude("test/llm/llm_test.rb")
+    .exclude("test/llm/gemma_test.rb")
+    .exclude("test/llm/mistral_test.rb")
+    .exclude("test/llm/llama_test.rb")
+    .exclude("test/llm/phi_test.rb")
+    .exclude("test/llm/qwen_test.rb")
 end
 
 spec = Bundler.load_gemspec("candle.gemspec")
@@ -63,6 +70,44 @@ task "test:device:benchmark" => :compile do
   Rake::Task["test:benchmark"].invoke
 end
 
+desc "Run LLM tests for specific models"
+namespace :test do
+  namespace :llm do
+    desc "Run tests for Gemma models"
+    task :gemma => :compile do
+      ruby "-Itest", "test/llm/gemma_test.rb"
+    end
+    
+    desc "Run tests for Phi models"
+    task :phi => :compile do
+      ruby "-Itest", "test/llm/phi_test.rb"
+    end
+    
+    desc "Run tests for Qwen models"
+    task :qwen => :compile do
+      ruby "-Itest", "test/llm/qwen_test.rb"
+    end
+    
+    desc "Run tests for Mistral models"
+    task :mistral => :compile do
+      ruby "-Itest", "test/llm/mistral_test.rb"
+    end
+    
+    desc "Run tests for Llama models"
+    task :llama => :compile do
+      ruby "-Itest", "test/llm/llama_test.rb"
+    end
+
+    desc "Run tests for TinyLlama models"
+    task :tinyllama => :compile do
+      ruby "-Itest", "test/llm/tinyllama_test.rb"
+    end
+
+    desc "Run all LLM tests (WARNING: downloads large models)"
+    task :all => [:gemma, :phi, :qwen, :mistral, :llama]
+  end
+end
+
 namespace :doc do
   task default: %i[rustdoc yard]
 

data/ext/candle/Cargo.toml

@@ -3,6 +3,7 @@ name = "candle"
 version = "0.1.0"
 edition = "2021"
 build = "build.rs"
+rust-version = "1.85"
 
 [lib]
 crate-type = ["cdylib"]
@@ -20,6 +21,7 @@ serde_json = "1.0"
 serde = { version = "1.0", features = ["derive"] }
 tokio = { version = "1.45", features = ["rt", "macros"] }
 rand = "0.8"
+outlines-core = "0.2"
 
 [features]
 default = []

data/ext/candle/src/lib.rs

@@ -7,6 +7,7 @@ pub mod llm;
 pub mod ner;
 pub mod reranker;
 pub mod ruby;
+pub mod structured;
 pub mod tokenizer;
 
 // Configuration detection from build.rs
@@ -49,6 +50,7 @@ fn init(ruby: &Ruby) -> Result<()> {
     ruby::device::init(rb_candle)?;
     ruby::tensor::init(rb_candle)?;
     ruby::tokenizer::init(rb_candle)?;
+    ruby::structured::init_structured(rb_candle)?;
     candle_utils(rb_candle)?;
 
     Ok(())

data/ext/candle/src/llm/constrained_generation_test.rs

@@ -0,0 +1,123 @@
+#[cfg(test)]
+mod constrained_generation_tests {
+    use super::super::*;
+    use crate::structured::{VocabularyAdapter, SchemaProcessor};
+    use crate::tokenizer::{TokenizerWrapper, loader::TokenizerLoader};
+    
+    #[tokio::test]
+    async fn test_constrained_vs_unconstrained_generation() {
+        // This test demonstrates the difference between constrained and unconstrained generation
+        
+        // Load a tokenizer for testing
+        if let Ok(tokenizer) = TokenizerLoader::from_hf_hub("bert-base-uncased", None).await {
+            let wrapper = TokenizerWrapper::new(tokenizer);
+            
+            // Create vocabulary adapter
+            let vocabulary = VocabularyAdapter::from_tokenizer(&wrapper)
+                .expect("Should create vocabulary");
+            
+            // Create schema processor
+            let processor = SchemaProcessor::new();
+            
+            // Define a simple JSON schema for a yes/no response
+            let schema = r#"{
+                "type": "object",
+                "properties": {
+                    "answer": {
+                        "type": "string",
+                        "enum": ["yes", "no"]
+                    }
+                },
+                "required": ["answer"]
+            }"#;
+            
+            // Process schema into Index
+            let index = processor.process_schema(schema, &vocabulary)
+                .expect("Should process schema");
+            
+            // Test configuration with constraint
+            let mut config_with_constraint = GenerationConfig::default();
+            config_with_constraint.constraint = Some(index.clone());
+            config_with_constraint.max_length = 50;
+            
+            // Test configuration without constraint
+            let config_without_constraint = GenerationConfig::default();
+            
+            // Create text generation instances
+            let mut gen_constrained = TextGeneration::from_config(&config_with_constraint);
+            let mut gen_unconstrained = TextGeneration::from_config(&config_without_constraint);
+            
+            // Set EOS token
+            gen_constrained.set_eos_token_id(102); // BERT's [SEP] token
+            gen_unconstrained.set_eos_token_id(102);
+            
+            // Constraints are set internally - we can't directly verify them
+            // but we can test their effects in actual generation
+        }
+    }
+    
+    #[test]
+    fn test_constraint_configuration() {
+        // Test that we can create a TextGeneration with constraints
+        let config = GenerationConfig::default();
+        let _text_gen = TextGeneration::from_config(&config);
+        
+        // Test that we can create a TextGeneration from config
+        // Constraints are private implementation details
+    }
+    
+    #[test]
+    fn test_repetition_penalty() {
+        use candle_core::{Tensor, Device};
+        
+        let device = Device::Cpu;
+        let vocab_size = 10;
+        
+        // Create logits with some positive and negative values
+        let logits_vec: Vec<f32> = vec![1.0, -1.0, 2.0, -2.0, 0.0, 3.0, -3.0, 1.5, -1.5, 0.5];
+        let mut logits = Tensor::from_vec(logits_vec.clone(), vocab_size, &device).unwrap();
+        
+        // Create text generation with some tokens
+        let mut text_gen = TextGeneration::new(42, Some(1.0), None, None, 1.0, 64);
+        text_gen.push_token(0); // Token that had logit 1.0
+        text_gen.push_token(2); // Token that had logit 2.0
+        text_gen.push_token(5); // Token that had logit 3.0
+        
+        // Apply repetition penalty
+        text_gen.apply_repetition_penalty(&mut logits, 1.5, 10).unwrap();
+        
+        let penalized = logits.to_vec1::<f32>().unwrap();
+        
+        // Check that tokens in context were penalized
+        assert!(penalized[0] < logits_vec[0], "Positive logit should be reduced");
+        assert!(penalized[2] < logits_vec[2], "Positive logit should be reduced");
+        assert!(penalized[5] < logits_vec[5], "Positive logit should be reduced");
+        
+        // Check that other tokens remain unchanged
+        assert_eq!(penalized[1], logits_vec[1], "Unsampled token should be unchanged");
+        assert_eq!(penalized[3], logits_vec[3], "Unsampled token should be unchanged");
+    }
+    
+    #[test]
+    fn test_stop_conditions() {
+        let mut text_gen = TextGeneration::new(42, Some(1.0), None, None, 1.0, 64);
+        text_gen.set_eos_token_id(50256); // Common EOS token
+        
+        // Test max length stop
+        for i in 0..10 {
+            text_gen.push_token(i);
+        }
+        assert!(text_gen.should_stop(100, 10), "Should stop at max length");
+        assert!(!text_gen.should_stop(100, 20), "Should not stop before max length");
+        
+        // Test EOS token stop
+        assert!(text_gen.should_stop(50256, 100), "Should stop at EOS token");
+        assert!(!text_gen.should_stop(123, 100), "Should not stop at non-EOS token");
+        
+        // Test stop sequences
+        let stop_seqs = vec!["STOP".to_string(), "END".to_string()];
+        assert!(text_gen.check_stop_sequences("This is the STOP", &stop_seqs), "Should detect stop sequence");
+        assert!(text_gen.check_stop_sequences("The END", &stop_seqs), "Should detect stop sequence");
+        assert!(!text_gen.check_stop_sequences("Continue", &stop_seqs), "Should not detect stop sequence");
+    }
+}