add tests

src/autora/doc/pipelines/main.py

@@ -1,5 +1,6 @@
 import logging
 from timeit import default_timer as timer
+from typing import List
 
 import jsonlines
 import mlflow
@@ -18,7 +19,9 @@
 
 
 @app.command()
-def predict(data_file: str, model_path: str, sys_id: SystemPrompts, instruc_id: InstructionPrompts) -> None:
+def predict(
+    data_file: str, model_path: str, sys_id: SystemPrompts, instruc_id: InstructionPrompts
+) -> List[str]:
     run = mlflow.active_run()
 
     sys_prompt = SYS[sys_id]
@@ -51,6 +54,7 @@ def predict(data_file: str, model_path: str, sys_id: SystemPrompts, instruc_id:
         total_tokens = sum([len(token) for token in tokens])
         mlflow.log_metric("total_tokens", total_tokens)
         mlflow.log_metric("tokens/sec", total_tokens / pred_time)
+        return predictions
 
 
 @app.command()

src/autora/doc/runtime/prompts.py

@@ -29,7 +29,6 @@ class SystemPrompts(Enum):
 
 
 class InstructionPrompts(Enum):
-    SYS_1 = "SYS_1"
     INSTR_SWEETP_1 = "INSTR_SWEETP_1"
 
 

tests/test_main.py

@@ -0,0 +1,15 @@
+from pathlib import Path
+
+from autora.doc.pipelines.main import predict
+from autora.doc.runtime.prompts import InstructionPrompts, SystemPrompts
+
+# dummy HF model for testing
+TEST_HF_MODEL = "hf-internal-testing/tiny-random-FalconForCausalLM"
+
+
+def test_predict() -> None:
+    data = Path(__file__).parent.joinpath("../data/data.jsonl").resolve()
+    outputs = predict(str(data), TEST_HF_MODEL, SystemPrompts.SYS_1, InstructionPrompts.INSTR_SWEETP_1)
+    assert len(outputs) == 3, "Expected 3 outputs"
+    for output in outputs:
+        assert len(output) > 0, "Expected non-empty output"