Coverage for src/trapi_predict_kit/types.py: 100%

27 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-07-25 21:14 +0200

1from typing import List, Optional 

2 

3from pydantic import BaseModel 

4 

5 

6class PredictHit(BaseModel): 

7 id: str 

8 type: str 

9 score: float 

10 label: Optional[str] 

11 

12 

13class PredictOutput(BaseModel): 

14 hits: List[PredictHit] 

15 count: int 

16 # input_id: str 

17 # input_type: str 

18 

19 

20class PredictOptions(BaseModel): 

21 model_id: Optional[str] = "openpredict_baseline" 

22 min_score: Optional[float] = None 

23 max_score: Optional[float] = None 

24 n_results: Optional[int] = None 

25 types: Optional[List[str]] = None 

26 

27 class Config: 

28 arbitrary_types_allowed = True 

29 

30 

31class TrainingOutput(BaseModel): 

32 # All scores are floats between 0 and 1 

33 precision: float 

34 recall: float 

35 accuracy: float 

36 roc_auc: float 

37 f1: float 

38 average_precision: float 

39 # elapsed_time: datetime 

40 

41 class Config: 

42 arbitrary_types_allowed = True