diff --git a/crawl4ai/model_loader.py b/crawl4ai/model_loader.py index aa80f673d..68f115d2d 100644 --- a/crawl4ai/model_loader.py +++ b/crawl4ai/model_loader.py @@ -75,8 +75,8 @@ def get_home_folder(): def load_bert_base_uncased(): from transformers import BertTokenizer, BertModel - tokenizer = BertTokenizer.from_pretrained("bert-base-uncased", resume_download=None) - model = BertModel.from_pretrained("bert-base-uncased", resume_download=None) + tokenizer = BertTokenizer.from_pretrained("bert-base-uncased") + model = BertModel.from_pretrained("bert-base-uncased") model.eval() model, device = set_model_device(model) return tokenizer, model @@ -94,8 +94,8 @@ def load_HF_embedding_model(model_name="BAAI/bge-small-en-v1.5") -> tuple: """ from transformers import AutoTokenizer, AutoModel - tokenizer = AutoTokenizer.from_pretrained(model_name, resume_download=None) - model = AutoModel.from_pretrained(model_name, resume_download=None) + tokenizer = AutoTokenizer.from_pretrained(model_name) + model = AutoModel.from_pretrained(model_name) model.eval() model, device = set_model_device(model) return tokenizer, model @@ -134,10 +134,8 @@ def load_text_multilabel_classifier(): # # return load_spacy_model(), torch.device("cpu") MODEL = "cardiffnlp/tweet-topic-21-multi" - tokenizer = AutoTokenizer.from_pretrained(MODEL, resume_download=None) - model = AutoModelForSequenceClassification.from_pretrained( - MODEL, resume_download=None - ) + tokenizer = AutoTokenizer.from_pretrained(MODEL) + model = AutoModelForSequenceClassification.from_pretrained(MODEL) model.eval() model, device = set_model_device(model) class_mapping = model.config.id2label diff --git a/tests/unit/test_model_loader_transformers.py b/tests/unit/test_model_loader_transformers.py new file mode 100644 index 000000000..390b4eaa1 --- /dev/null +++ b/tests/unit/test_model_loader_transformers.py @@ -0,0 +1,87 @@ +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from crawl4ai import model_loader + + +class _FakePretrained: + from_pretrained = Mock() + + +@pytest.fixture(autouse=True) +def reset_pretrained_mock(): + _FakePretrained.from_pretrained.reset_mock() + + +def _fake_model(): + model = Mock() + model.config.id2label = {} + return model + + +def test_bert_loader_uses_current_from_pretrained_api(monkeypatch): + tokenizer = object() + model = _fake_model() + _FakePretrained.from_pretrained.side_effect = [tokenizer, model] + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace(BertTokenizer=_FakePretrained, BertModel=_FakePretrained), + ) + monkeypatch.setattr(model_loader, "set_model_device", lambda value: (value, "cpu")) + + loaded_tokenizer, loaded_model = model_loader.load_bert_base_uncased() + + assert (loaded_tokenizer, loaded_model) == (tokenizer, model) + assert _FakePretrained.from_pretrained.call_args_list == [ + (("bert-base-uncased",), {}), + (("bert-base-uncased",), {}), + ] + + +def test_embedding_loader_uses_current_from_pretrained_api(monkeypatch): + tokenizer = object() + model = _fake_model() + _FakePretrained.from_pretrained.side_effect = [tokenizer, model] + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace(AutoTokenizer=_FakePretrained, AutoModel=_FakePretrained), + ) + monkeypatch.setattr(model_loader, "set_model_device", lambda value: (value, "cpu")) + + loaded_tokenizer, loaded_model = model_loader.load_HF_embedding_model("example/model") + + assert (loaded_tokenizer, loaded_model) == (tokenizer, model) + assert _FakePretrained.from_pretrained.call_args_list == [ + (("example/model",), {}), + (("example/model",), {}), + ] + + +def test_multilabel_loader_uses_current_from_pretrained_api(monkeypatch): + tokenizer = object() + model = _fake_model() + _FakePretrained.from_pretrained.side_effect = [tokenizer, model] + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace( + AutoTokenizer=_FakePretrained, + AutoModelForSequenceClassification=_FakePretrained, + ), + ) + monkeypatch.setitem(sys.modules, "torch", SimpleNamespace()) + monkeypatch.setattr(model_loader, "set_model_device", lambda value: (value, "cpu")) + + classifier, device = model_loader.load_text_multilabel_classifier() + + assert callable(classifier) + assert device == "cpu" + assert _FakePretrained.from_pretrained.call_args_list == [ + (("cardiffnlp/tweet-topic-21-multi",), {}), + (("cardiffnlp/tweet-topic-21-multi",), {}), + ]