Add basic LlamaState tests

abetlen · abetlen · commit e552955217b4 · 2024-09-18T21:11:05.000-04:00
diff --git a/tests/test_llama.py b/tests/test_llama.py
@@ -171,3 +171,8 @@ def logit_processor_func(input_ids, logits):
         logits_processor=logit_processors
     )
     assert output["choices"][0]["text"].lower().startswith("rot")
+
+    # TODO: Add test better test here
+    state = model.save_state()
+    model.load_state(state)
+

Original file line number	Diff line number	Diff line change
`@@ -171,3 +171,8 @@ def logit_processor_func(input_ids, logits):`
`171`	`171`	`logits_processor=logit_processors`
`172`	`172`	`)`
`173`	`173`	`assert output["choices"][0]["text"].lower().startswith("rot")`
	`174`	`+`
	`175`	`+ # TODO: Add test better test here`
	`176`	`+ state = model.save_state()`
	`177`	`+ model.load_state(state)`
	`178`	`+`