|
17 | 17 | import tempfile
|
18 | 18 | import unittest
|
19 | 19 |
|
20 |
| -import pytest |
21 | 20 | from executorch.extension.pybindings.portable_lib import ExecuTorchModule
|
22 |
| -from transformers.testing_utils import slow |
23 | 21 |
|
24 | 22 | from optimum.executorch import ExecuTorchModelForCausalLM
|
| 23 | +from optimum.exporters.executorch import main_export |
25 | 24 |
|
26 | 25 |
|
27 | 26 | class ExecuTorchModelIntegrationTest(unittest.TestCase):
|
28 | 27 | def __init__(self, *args, **kwargs):
|
29 | 28 | super().__init__(*args, **kwargs)
|
30 | 29 |
|
31 |
| - @slow |
32 |
| - @pytest.mark.run_slow |
33 | 30 | def test_load_model_from_hub(self):
|
34 |
| - model = ExecuTorchModelForCausalLM.from_pretrained("NousResearch/Llama-3.2-1B", export=True, recipe="xnnpack") |
| 31 | + model_id = "optimum-internal-testing/tiny-random-llama" |
| 32 | + |
| 33 | + model = ExecuTorchModelForCausalLM.from_pretrained(model_id, export=True, recipe="xnnpack") |
35 | 34 | self.assertIsInstance(model, ExecuTorchModelForCausalLM)
|
36 | 35 | self.assertIsInstance(model.model, ExecuTorchModule)
|
37 | 36 |
|
38 |
| - @slow |
39 |
| - @pytest.mark.run_slow |
40 |
| - def test_load_model_from_local_path(self): |
41 |
| - from optimum.exporters.executorch import main_export |
| 37 | + def test_load_et_model_from_hub(self): |
| 38 | + model_id = "optimum-internal-testing/tiny-random-llama" |
| 39 | + |
| 40 | + model = ExecuTorchModelForCausalLM.from_pretrained( |
| 41 | + model_id, export=False, revision="executorch", recipe="xnnpack" |
| 42 | + ) |
| 43 | + self.assertIsInstance(model, ExecuTorchModelForCausalLM) |
| 44 | + self.assertIsInstance(model.model, ExecuTorchModule) |
42 | 45 |
|
43 |
| - model_id = "NousResearch/Llama-3.2-1B" |
44 |
| - task = "text-generation" |
| 46 | + def test_load_model_from_local_path(self): |
| 47 | + model_id = "optimum-internal-testing/tiny-random-llama" |
45 | 48 | recipe = "xnnpack"
|
46 | 49 |
|
47 | 50 | with tempfile.TemporaryDirectory() as tempdir:
|
48 | 51 | # Export to a local dir
|
49 | 52 | main_export(
|
50 | 53 | model_name_or_path=model_id,
|
51 |
| - task=task, |
52 | 54 | recipe=recipe,
|
53 | 55 | output_dir=tempdir,
|
54 | 56 | )
|
|
0 commit comments