Add Qwen2.5, Gemma-1, Phi-2

tests/models/gemma_2b_it/test_gemma_2b_it.py

@@ -0,0 +1,44 @@
+import torch
+import pytest
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from tests.utils import ModelTester
+
+
+class ThisTester(ModelTester):
+    def _load_model(self):
+        model_name = "google/gemma-1.1-2b-it"
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name, padding_side="left", torch_dtype=torch.bfloat16
+        )
+        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+        return model.generate
+
+    def _load_inputs(self):
+        prompt = "This is a test prompt."
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
+        arguments = {
+            "input_ids": input_ids,
+            "do_sample": True,
+            "temperature": 0.9,
+            "max_new_tokens": 5,
+        }
+        return arguments
+
+    def set_model_eval(self, model):
+        return model
+
+
+@pytest.mark.parametrize("mode", ["eval"])
+@pytest.mark.compilation_xfail
+def test_gemma(record_property, mode):
+    model_name = "Gemma-1.1-2b-it"
+    record_property("model_name", model_name)
+    record_property("mode", mode)
+
+    tester = ThisTester(model_name, mode)
+    results = tester.test_model()
+    if mode == "eval":
+        gen_text = tester.tokenizer.batch_decode(results)[0]
+        print(f"Generated Text:\n{gen_text}")
+
+    record_property("torch_ttnn", (tester, results))

tests/models/phi_2/test_phi.py

@@ -0,0 +1,44 @@
+import torch
+import pytest
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from tests.utils import ModelTester
+
+
+class ThisTester(ModelTester):
+    def _load_model(self):
+        model_name = "microsoft/phi-2"
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name, padding_side="left", torch_dtype=torch.bfloat16
+        )
+        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+        return model.generate
+
+    def _load_inputs(self):
+        prompt = "This is a test prompt."
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
+        arguments = {
+            "input_ids": input_ids,
+            "do_sample": True,
+            "temperature": 0.9,
+            "max_new_tokens": 5,
+        }
+        return arguments
+
+    def set_model_eval(self, model):
+        return model
+
+
+@pytest.mark.parametrize("mode", ["eval"])
+@pytest.mark.compilation_xfail
+def test_phi_2(record_property, mode):
+    model_name = "Phi-2"
+    record_property("model_name", model_name)
+    record_property("mode", mode)
+
+    tester = ThisTester(model_name, mode)
+    results = tester.test_model()
+    if mode == "eval":
+        gen_text = tester.tokenizer.batch_decode(results)[0]
+        print(f"Generated Text:\n{gen_text}")
+
+    record_property("torch_ttnn", (tester, results))

tests/models/qwen2.5/test_qwen2.py

@@ -0,0 +1,44 @@
+import torch
+import pytest
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from tests.utils import ModelTester
+
+
+class ThisTester(ModelTester):
+    def _load_model(self):
+        model_name = "Qwen/Qwen2.5-0.5B-Instruct"
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name, padding_side="left", torch_dtype=torch.bfloat16
+        )
+        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+        return model.generate
+
+    def _load_inputs(self):
+        prompt = "This is a test prompt."
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
+        arguments = {
+            "input_ids": input_ids,
+            "do_sample": True,
+            "temperature": 0.9,
+            "max_new_tokens": 5,
+        }
+        return arguments
+
+    def set_model_eval(self, model):
+        return model
+
+
+@pytest.mark.parametrize("mode", ["eval"])
+@pytest.mark.compilation_xfail
+def test_qwen(record_property, mode):
+    model_name = "Qwen2.5"
+    record_property("model_name", model_name)
+    record_property("mode", mode)
+
+    tester = ThisTester(model_name, mode)
+    results = tester.test_model()
+    if mode == "eval":
+        gen_text = tester.tokenizer.batch_decode(results)[0]
+        print(f"Generated Text:\n{gen_text}")
+
+    record_property("torch_ttnn", (tester, results))

tests/models/qwen2_vl/test_qwen.py

@@ -0,0 +1,68 @@
+import torch
+import pytest
+from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
+from PIL import Image
+import requests
+from torchvision import io
+from typing import Dict
+from tests.utils import ModelTester
+
+
+class ThisTester(ModelTester):
+    def _load_model(self):
+        model_name = "Qwen/Qwen2-VL-2B-Instruct"
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name, padding_side="left", torch_dtype=torch.bfloat16
+        )
+        model = Qwen2VLForConditionalGeneration.from_pretrained(
+            "Qwen/Qwen2-VL-2B-Instruct", torch_dtype=torch.bfloat16
+        )
+        self.processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
+
+        return model.generate
+
+    def _load_inputs(self):
+        # Image
+        url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
+        image = Image.open(requests.get(url, stream=True).raw)
+        conversation = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image",
+                    },
+                    {"type": "text", "text": "Describe this image."},
+                ],
+            }
+        ]
+        text_prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=True)
+        inputs = self.processor(
+            text=[text_prompt], images=[image], padding=True, return_tensors="pt"
+        )
+
+        arguments = inputs.update({
+            "do_sample": True,
+            "temperature": 0.9,
+            "max_new_tokens": 10,
+        })
+        return arguments
+
+    def set_model_eval(self, model):
+        return model
+
+
+@pytest.mark.parametrize("mode", ["eval"])
+@pytest.mark.skip(reason="This test requires transformers>=4.45.0")
+def test_qwen(record_property, mode):
+    model_name = "Qwen2-VL-2B"
+    record_property("model_name", model_name)
+    record_property("mode", mode)
+
+    tester = ThisTester(model_name, mode)
+    results = tester.test_model()
+    if mode == "eval":
+        gen_text = tester.processor.batch_decode(results, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+        print(f"Generated Text:\n{gen_text}")
+
+    record_property("torch_ttnn", (tester, results))