Merge pull request #325 from explosion/fix/model-init

Fix Llama 2 model init and Anthropic requests
explosion · Oct 13, 2023 · 89b5b65 · 89b5b65
2 parents 3a8b28e + 1240146
commit 89b5b65
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 12 deletions.
diff --git a/spacy_llm/models/hf/llama2.py b/spacy_llm/models/hf/llama2.py
@@ -31,7 +31,6 @@ def init_model(self) -> Any:
         return transformers.pipeline(
             "text-generation",
             model=self._name,
-            use_auth_token=True,
             return_full_text=False,
             **self._config_init,
         )
@@ -48,14 +47,7 @@ def __call__(self, prompts: Iterable[str]) -> Iterable[str]:  # type: ignore[ove
 
     @staticmethod
     def compile_default_configs() -> Tuple[Dict[str, Any], Dict[str, Any]]:
-        default_cfg_init, default_cfg_run = HuggingFace.compile_default_configs()
-        return (
-            {
-                **default_cfg_init,
-                "trust_remote_code": True,
-            },
-            default_cfg_run,
-        )
+        return HuggingFace.compile_default_configs()
 
 
 @registry.llm_models("spacy.Llama2.v1")

diff --git a/spacy_llm/models/rest/anthropic/model.py b/spacy_llm/models/rest/anthropic/model.py
@@ -54,6 +54,7 @@ def __call__(self, prompts: Iterable[str]) -> Iterable[str]:
         headers = {
             **self._credentials,
             "model": self._name,
+            "anthropic_version": self._config.get("anthropic_version", "2023-06-01"),
             "Content-Type": "application/json",
         }
 

diff --git a/spacy_llm/tests/models/test_llama2.py b/spacy_llm/tests/models/test_llama2.py
@@ -37,7 +37,7 @@
 """
 
 
-@pytest.mark.skip(reason="CI runner needs more GPU memory")
+# @pytest.mark.skip(reason="CI runner needs more GPU memory")
 @pytest.mark.gpu
 @pytest.mark.skipif(not has_torch_cuda_gpu, reason="needs GPU & CUDA")
 def test_init():
@@ -52,7 +52,7 @@ def test_init():
     )
 
 
-@pytest.mark.skip(reason="CI runner needs more GPU memory")
+# @pytest.mark.skip(reason="CI runner needs more GPU memory")
 @pytest.mark.gpu
 @pytest.mark.skipif(not has_torch_cuda_gpu, reason="needs GPU & CUDA")
 def test_init_from_config():
@@ -62,7 +62,7 @@ def test_init_from_config():
     torch.cuda.empty_cache()
 
 
-@pytest.mark.skip(reason="CI runner needs more GPU memory")
+# @pytest.mark.skip(reason="CI runner needs more GPU memory")
 @pytest.mark.gpu
 @pytest.mark.skipif(not has_torch_cuda_gpu, reason="needs GPU & CUDA")
 def test_invalid_model():