Phi3 tests: fix typing for Python 3.8 (#32388)

fix phi
This commit is contained in:
Raushan Turganbay 2024-08-05 11:58:42 +05:00 committed by GitHub
parent 05ae3a300d
commit 3bb646a54f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -16,6 +16,7 @@
"""Testing suite for the PyTorch Phi-3 model."""
import unittest
from typing import List
from parameterized import parameterized
@ -69,7 +70,7 @@ if is_torch_available():
).logits
@staticmethod
def generate(model: Phi3ForCausalLM, prompt_tokens: torch.LongTensor, max_seq_len: int) -> list[int]:
def generate(model: Phi3ForCausalLM, prompt_tokens: torch.LongTensor, max_seq_len: int) -> List[int]:
model = Phi3MiniWithStaticCache(model, 1, max_seq_len + prompt_tokens.shape[-1])
response_tokens = []