Expose MockAsyncLM in load_model_by_name (#42)

* Expose MockAsyncLM in load_model_by_name

* Update test_llm.py

Author: benlebrun

genlm/backend/llm/__init__.py

@@ -10,7 +10,7 @@ def load_model_by_name(name, backend=None, llm_opts=None):
 
     Args:
         name (str): Hugging Face model name (e.g. "gpt2", "meta-llama/Llama-3.2-1B-Instruct")
-        backend (str, optional): Backend to use for inference. Can be "vllm" or "hf".
+        backend (str, optional): Backend to use for inference. Can be "vllm", "hf" or "mock".
             If None, defaults to "vllm" if CUDA is available, otherwise "hf".
         llm_opts (dict, optional): Additional options to pass to the backend constructor.
             See AsyncVirtualLM and AsyncTransformer documentation for details.
@@ -31,6 +31,8 @@ def load_model_by_name(name, backend=None, llm_opts=None):
         return AsyncVirtualLM.from_name(name, **llm_opts)
     elif backend == "hf":
         return AsyncTransformer.from_name(name, **llm_opts)
+    elif backend == "mock":
+        return MockAsyncLM.from_name(name, **llm_opts)
     else:
         raise ValueError(f"Invalid backend: {backend}")
 

tests/test_llm.py

@@ -226,6 +226,10 @@ async def test_mock_async_llm():
     mock_async_llm.clear_cache()  # no-op
 
 
+def test_load_model_by_name_mock():
+    load_model_by_name("gpt2", backend="mock")
+
+
 def test_load_model_by_name_error():
     with pytest.raises(ValueError):
         load_model_by_name("gpt2", backend="invalid")