Add actual logit processor test

abetlen · abetlen · commit 334ca35b715a · 2024-09-18T19:59:40.000-04:00
diff --git a/tests/test_llama.py b/tests/test_llama.py
@@ -15,6 +15,10 @@
 MODEL = "./vendor/llama.cpp/models/ggml-vocab-llama-spm.gguf"
 
 
+def test_llama_cpp_version():
+    assert llama_cpp.__version__
+
+
 def test_llama_cpp_tokenization():
     llama = llama_cpp.Llama(model_path=MODEL, vocab_only=True, verbose=False)
 
@@ -52,10 +56,6 @@ def test_llama_cpp_tokenization():
     assert text == llama.detokenize(tokens)
 
 
-def test_llama_cpp_version():
-    assert llama_cpp.__version__
-
-
 @pytest.fixture
 def llama_cpp_model_path():
     repo_id = "Qwen/Qwen2-0.5B-Instruct-GGUF"
@@ -150,8 +150,12 @@ def test_real_llama(llama_cpp_model_path):
     )
     assert output["choices"][0]["text"] == "true"
 
+    suffix = b"rot"
+    tokens = model.tokenize(suffix, add_bos=True, special=True)
     def logit_processor_func(input_ids, logits):
-        return logits * 1
+        for token in tokens:
+            logits[token] *= 1000
+        return logits
 
     logit_processors = llama_cpp.LogitsProcessorList(
         [logit_processor_func]
@@ -166,4 +170,4 @@ def logit_processor_func(input_ids, logits):
         seed=1337,
         logits_processor=logit_processors
     )
-    assert output["choices"][0]["text"].lower().startswith("is")
+    assert output["choices"][0]["text"].lower().startswith("rot")
