mirror of
https://github.com/deepset-ai/haystack.git
synced 2025-08-28 18:36:36 +00:00
25 lines
989 B
Python
25 lines
989 B
Python
|
from haystack import Document
|
||
|
|
from haystack.classifier.base import BaseClassifier
|
||
|
|
|
||
|
|
|
||
|
|
def test_classifier(classifier):
|
||
|
|
assert isinstance(classifier, BaseClassifier)
|
||
|
|
|
||
|
|
query = "not used at the moment"
|
||
|
|
docs = [
|
||
|
|
Document(
|
||
|
|
text="""Fragen und Antworten - Bitte auf Themen beschränken welche einen Bezug zur Bahn aufweisen. Persönliche Unterhaltungen bitte per PN führen. Links bitte mit kurzer Erklärung zum verlinkten Inhalt versehen""",
|
||
|
|
meta={"name": "0"},
|
||
|
|
id="1",
|
||
|
|
),
|
||
|
|
Document(
|
||
|
|
text="""Ich liebe es wenn die Bahn selbstverschuldete unnötig lange Aufenthaltszeiten durch Verspätung wieder rausfährt.""",
|
||
|
|
meta={"name": "1"},
|
||
|
|
id="2",
|
||
|
|
),
|
||
|
|
]
|
||
|
|
results = classifier.predict(query=query, documents=docs)
|
||
|
|
expected_labels = ["neutral", "negative"]
|
||
|
|
for i, doc in enumerate(results):
|
||
|
|
assert doc.to_dict()["meta"]["classification"]["label"] == expected_labels[i]
|