File size: 1,931 Bytes
b91a6bd | 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 | import unittest
from unittest.mock import patch
import pandas as pd
import torch
from utils import pipeline
class FakeEmbeddingModel:
def __init__(self, vectors):
self.vectors = vectors
self.calls = []
def __call__(self, text, pooling='mean'):
self.calls.append((text, pooling))
return torch.tensor(self.vectors[text])
class PipelineFactoryTest(unittest.TestCase):
def test_originality_uses_embedding_factory(self):
model = FakeEmbeddingModel(
{
'prompt': [1.0, 0.0],
'response': [0.0, 1.0],
}
)
df = pd.DataFrame({'prompt': ['prompt'], 'response': ['response']})
with patch.object(pipeline, 'get_embedding_model', return_value=model) as factory:
result = pipeline.p0_originality(df, 'fake-model', 'mean')
factory.assert_called_once_with('fake-model')
self.assertAlmostEqual(result.loc[0, 'originality'], 1.0)
self.assertEqual(model.calls, [('prompt', 'mean'), ('response', 'mean')])
def test_flexibility_uses_embedding_factory(self):
model = FakeEmbeddingModel(
{
'p': [1.0, 0.0],
'a': [1.0, 0.0],
'b': [0.0, 1.0],
}
)
df = pd.DataFrame(
{
'id': [1, 1, 1],
'prompt': ['p', 'p', 'p'],
'response': ['a', 'b', 'a'],
}
)
with patch.object(pipeline, 'get_embedding_model', return_value=model) as factory:
result = pipeline.p1_flexibility(df, 'fake-model', 'cls')
factory.assert_called_once_with('fake-model')
self.assertEqual(len(result), 1)
self.assertAlmostEqual(result.loc[0, 'flexibility'], 2.0)
self.assertEqual(model.calls, [('a', 'cls'), ('b', 'cls'), ('a', 'cls')])
if __name__ == '__main__':
unittest.main()
|