| """Task classification on the sample pairs (rules; embeddings when cached).""" |
|
|
| import samples |
| from rubriq.classify import classify, split_instruction_and_source |
| from rubriq.schemas import TaskType |
|
|
|
|
| def test_summarization_detected(): |
| c = classify(samples.SUMM_PROMPT, samples.SUMM_GOOD) |
| assert c.task_type == TaskType.SUMMARIZATION |
| assert c.has_source_text |
|
|
|
|
| def test_code_detected(): |
| c = classify(samples.CODE_PROMPT, samples.CODE_GOOD) |
| assert c.task_type == TaskType.CODE_GENERATION |
|
|
|
|
| def test_extraction_detected(): |
| c = classify(samples.EXTRACT_PROMPT, samples.EXTRACT_GOOD) |
| assert c.task_type == TaskType.EXTRACTION |
| assert c.has_source_text |
|
|
|
|
| def test_creative_detected(): |
| c = classify(samples.CREATIVE_PROMPT, samples.CREATIVE_CLICHED) |
| assert c.task_type == TaskType.CREATIVE_WRITING |
|
|
|
|
| def test_open_qa_detected(): |
| c = classify(samples.QA_OPEN_PROMPT, samples.QA_OPEN_RISKY) |
| assert c.task_type == TaskType.QA_OPEN |
| assert not c.has_source_text |
|
|
|
|
| def test_grounded_qa_upgrade(): |
| c = classify(samples.QA_MULTI_PROMPT, samples.QA_MULTI_PARTIAL) |
| assert c.task_type == TaskType.QA_GROUNDED |
| assert c.has_source_text |
|
|
|
|
| def test_source_split_keeps_instruction(): |
| instruction, source = split_instruction_and_source(samples.SUMM_PROMPT) |
| assert "at most 60 words" in instruction |
| assert source is not None and "Solheim" in source |
|
|
