Spaces:
Running
Running
File size: 3,231 Bytes
dea4a7a |
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 |
import json
import unittest
from unittest.mock import MagicMock, patch
from langchain.chat_models import ChatOpenAI
from meta_prompt.sample_generator import TaskDescriptionGenerator
class TestTaskDescriptionGenerator(unittest.TestCase):
def setUp(self):
self.model = ChatOpenAI(model="llama3-70b-8192", temperature=1.0, max_retries=3)
self.generator = TaskDescriptionGenerator(self.model)
@patch.object(ChatOpenAI, "call")
def test_generate_description(self, mock_call):
mock_call.return_value = "Task Description: Describe a cat."
input_json = json.dumps({"input": "A cat", "output": "A furry animal"})
description = self.generator.generate_description(input_json)
self.assertEqual(description, "Task Description: Describe a cat.")
@patch.object(ChatOpenAI, "call")
def test_analyze_input(self, mock_call):
mock_call.return_value = "Input Analysis: The input is an animal."
description = "Task Description: Describe a cat."
input_analysis = self.generator.analyze_input(description)
self.assertEqual(input_analysis, "Input Analysis: The input is an animal.")
@patch.object(ChatOpenAI, "call")
def test_generate_briefs(self, mock_call):
mock_call.return_value = '{"new_example_briefs": [{"example_brief": "Brief 1"}, {"example_brief": "Brief 2"}]}'
description = "Task Description: Describe a cat."
input_analysis = "Input Analysis: The input is an animal."
generating_batch_size = 2
briefs = self.generator.generate_briefs(description, input_analysis, generating_batch_size)
self.assertEqual(briefs, {"new_example_briefs": [{"example_brief": "Brief 1"}, {"example_brief": "Brief 2"}]})
@patch.object(ChatOpenAI, "call")
def test_generate_examples_from_briefs(self, mock_call):
mock_call.return_value = '{"examples": [{"input": "Input 1", "output": "Output 1"}, {"input": "Input 2", "output": "Output 2"}]}'
description = "Task Description: Describe a cat."
new_example_briefs = {"new_example_briefs": [{"example_brief": "Brief 1"}, {"example_brief": "Brief 2"}]}
raw_example = json.dumps({"input": "A cat", "output": "A furry animal"})
generating_batch_size = 2
examples = self.generator.generate_examples_from_briefs(description, new_example_briefs, raw_example, generating_batch_size)
self.assertEqual(examples, {"examples": [{"input": "Input 1", "output": "Output 1"}, {"input": "Input 2", "output": "Output 2"}]})
@patch.object(ChatOpenAI, "call")
def test_generate_examples(self, mock_call):
mock_call.return_value = '{"examples": [{"input": "Input 1", "output": "Output 1"}, {"input": "Input 2", "output": "Output 2"}]}'
description = "Task Description: Describe a cat."
raw_example = json.dumps({"input": "A cat", "output": "A furry animal"})
generating_batch_size = 2
examples = self.generator.generate_examples(description, raw_example, generating_batch_size)
self.assertEqual(examples, {"examples": [{"input": "Input 1", "output": "Output 1"}, {"input": "Input 2", "output": "Output 2"}]})
if __name__ == '__main__':
unittest.main() |