File size: 1,399 Bytes
4b5b8fc | 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 | """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
|