Coverage for rdgai/classification.py: 100.00%

1from pathlib import Path 

2from langchain_core.output_parsers import StrOutputParser 

3from langchain_core.language_models.llms import LLM 

4import llmloader 

5from rich.console import Console 

6from rich.progress import track 

7 

8from .prompts import build_template 

9from .parsers import CategoryParser 

10from .apparatus import Doc, Pair 

11 

12 

13DEFAULT_MODEL_ID = "gpt-4o" 

14 

15 

16def classify_pair( 

17 doc:Doc, 

18 pair:Pair, 

19 llm:LLM, 

20 output:Path, 

21 verbose:bool=False, 

22 prompt_only:bool=False, 

23 examples:int=10, 

24 console:Console|None=None, 

25 examples_doc:Doc|None=None, 

26): 

27 """ 

28 Classifies relations for a pair of readings. 

29 """ 

30 assert isinstance(doc, Doc), f"Expected Doc, got {type(doc)}" 

31 

32 console = console or Console() 

33 

34 template = build_template(pair, examples=examples, examples_doc=examples_doc) 

35 if verbose or prompt_only: 

36 template.pretty_print() 

37 if prompt_only: 

38 return 

39 

40 chain = template | llm | StrOutputParser() | CategoryParser(doc.relation_types.keys()) 

41 

42 assert isinstance(output, Path), f"Expected Path, got {type(output)}" 

43 doc.write(output) 

44 

45 category, description = chain.invoke({}) 

46 

47 console.print() 

48 pair.print(console) 

49 console.print(category, style="green bold") 

50 console.print(description, style="grey46") 

51 

52 relation_type = doc.relation_types.get(category, None) 

53 if relation_type is None: 

54 return 

55 

56 inverse_description = f"c.f. {pair.active} ➞ {pair.passive}" 

57 pair.add_type_with_inverse( 

58 relation_type, 

59 responsible="#rdgai", 

60 description=description, 

61 inverse_description=inverse_description, 

62 ) 

63 

64 doc.write(output) 

65 

66 

67def classify( 

68 doc:Doc, 

69 output:Path, 

70 pairs:list[Pair]|None=None, 

71 verbose:bool=False, 

72 api_key:str="", 

73 llm:str=DEFAULT_MODEL_ID, 

74 temperature:float=0.1, 

75 prompt_only:bool=False, 

76 examples:int=10, 

77 console:Console|None=None, 

78 examples_doc:Doc|None=None, 

79): 

80 """ 

81 Classifies relations in TEI documents. 

82 """ 

83 assert isinstance(doc, Doc), f"Expected Doc, got {type(doc)}" 

84 

85 console = console or Console() 

86 llm = llmloader.load(model=llm, api_key=api_key, temperature=temperature) 

87 

88 pairs = pairs or doc.get_unclassified_pairs(redundant=False) 

89 for pair in track(pairs): 

90 classify_pair( 

91 doc, 

92 pair, 

93 llm, 

94 output, 

95 verbose=verbose, 

96 prompt_only=prompt_only, 

97 examples=examples, 

98 console=console, 

99 examples_doc=examples_doc, 

100 ) 

101 

102