Spaces:
Running
Running
File size: 2,090 Bytes
1965493 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
"""
Common test utilities and mocks for unit tests.
"""
from unittest.mock import MagicMock
class MockBertTokenizer:
"""
A mock for transformers.BertTokenizer for testing purposes.
"""
def __init__(self, *args, **kwargs):
"""Initialize the mock tokenizer."""
self.vocab = {"[PAD]": 0, "[UNK]": 1}
self.model_max_length = 512
def encode(self, text, add_special_tokens=True, truncation=True, max_length=None):
"""
Mock encode method to convert text to token IDs.
"""
# Return some dummy token IDs
return [1, 2, 3]
def decode(self, token_ids, skip_special_tokens=True):
"""
Mock decode method to convert token IDs back to text.
"""
# Return dummy text
return 'decoded text'
def __call__(self, text, padding=True, truncation=True, max_length=None, return_tensors=None):
"""
Mock call method to simulate tokenization.
"""
return {
'input_ids': [[1, 2, 3]],
'attention_mask': [[1, 1, 1]]
}
def patch_bert_tokenizer():
"""
Returns a mock for transformers.BertTokenizer
"""
mock_tokenizer = MagicMock()
mock_tokenizer.from_pretrained = MagicMock(return_value=MockBertTokenizer())
return mock_tokenizer
def get_mock_llm_response():
"""
Returns a mock LLM response for testing
"""
return '''
{
"title": "Test Presentation",
"slides": [
{
"title": "Test Slide 1",
"content": "Test content",
"layout": "text_only"
}
]
}
'''
class MockStreamResponse:
"""
A mock class to simulate streaming responses from an LLM.
"""
def __init__(self, content):
self.content = content
def __iter__(self):
yield self
def get_mock_llm():
"""
Returns a mock LLM instance for testing
"""
mock_llm = MagicMock()
mock_llm.stream.return_value = [MockStreamResponse(get_mock_llm_response())]
return mock_llm
|