Coverage for rdgai/main.py: 100.00%

1from pathlib import Path 

2import typer 

3from rich.console import Console 

4import pandas as pd 

5 

6from .apparatus import Doc 

7from .export import export_variants_to_excel, import_classifications_from_dataframe 

8from .classification import classify as classify_fn 

9from .evaluation import evaluate_docs 

10from .classification import DEFAULT_MODEL_ID 

11from .validation import validate as validate_fn 

12from .prompts import build_preamble 

13 

14console = Console() 

15error_console = Console(stderr=True, style="bold red") 

16 

17 

18app = typer.Typer(pretty_exceptions_enable=False) 

19 

20 

21 

22def get_output_path(doc:Path, output:Path, inplace:bool) -> Path: 

23 """ Checks if the output path should be replaced with the input doc. """ 

24 if output and inplace: 

25 raise typer.BadParameter("You cannot use both an output path and --inplace/-i at the same time.") 

26 if not output and not inplace: 

27 raise typer.BadParameter("You must provide either an output path or use --inplace/-i.") 

28 

29 if inplace: 

30 output = doc 

31 

32 if isinstance(output, Doc): 

33 output = output.path 

34 

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

36 

37 return output 

38 

39 

40@app.command() 

41def classify( 

42 doc:Path=typer.Argument(..., help="The path to the TEI XML document to classify."), 

43 output:Path=typer.Argument(None, help="The path to the output TEI XML file."), 

44 inplace: bool = typer.Option(False, "--inplace", "-i", help="Overwrite the input file."), 

45 verbose:bool=typer.Option(False, help="Print verbose output."), 

46 api_key:str=typer.Option("", help="API key for the LLM."), 

47 llm:str=typer.Option(DEFAULT_MODEL_ID, help="ID of the language model to use."), 

48 temperature:float=typer.Option(0.1, help="Temperature for sampling from the language model."), 

49 prompt_only:bool=typer.Option(False, help="Only print the prompt and not classify."), 

50 examples:int=typer.Option(10, help="Number of examples to include in the prompt."), 

51 examples_doc:Path=typer.Option(None, help="The path to a TEI XML document to use for examples.") 

52): 

53 """ 

54 Classifies relations in TEI documents. 

55 """ 

56 doc = Doc(doc) 

57 output = get_output_path(doc, output, inplace) 

58 examples_doc = Doc(examples_doc) if examples_doc and Path(examples_doc).exists() else None 

59 

60 return classify_fn( 

61 doc=doc, 

62 output=output, 

63 verbose=verbose, 

64 api_key=api_key, 

65 llm=llm, 

66 temperature=temperature, 

67 prompt_only=prompt_only, 

68 examples=examples, 

69 console=console, 

70 examples_doc=examples_doc, 

71 ) 

72 

73 

74@app.command() 

75def classified_pairs( 

76 doc:Path=typer.Argument(..., help="The path to the TEI XML document with the classifications."), 

77): 

78 """ Print classified pairs in a document. """ 

79 doc = Doc(doc) 

80 doc.print_classified_pairs(console) 

81 

82 

83@app.command() 

84def html( 

85 doc:Path=typer.Argument(..., help="The path to the TEI XML document to render as HTML."), 

86 output:Path=typer.Argument(..., help="The path to the output HTML file."), 

87 all_apps:bool=typer.Option(False, help="Whether or not to use all variation unit `app` elements. By default it shows only non-redundant pairs of readings."), 

88): 

89 """ Renders the variation units of a TEI document as HTML. """ 

90 doc = Doc(doc) 

91 doc.render_html(output, all_apps=all_apps) 

92 

93 

94@app.command() 

95def gui( 

96 doc:Path=typer.Argument(..., help="The path to the TEI XML document to classify."), 

97 output:Path=typer.Argument(None, help="The path to the output TEI XML file."), 

98 inplace: bool = typer.Option(False, "--inplace", "-i", help="Overwrite the input file."), 

99 debug:bool=True, 

100 use_reloader:bool=False, 

101 all_apps:bool=typer.Option(False, help="Whether or not to use all variation unit `app` elements. By default it shows only non-redundant pairs of readings."), 

102): 

103 """ Starts a Flask app to view and classify a TEI document. """ 

104 output = get_output_path(doc, output, inplace) 

105 doc = Doc(doc) 

106 flask_app = doc.flask_app(output, all_apps=all_apps) 

107 flask_app.run(debug=debug, use_reloader=use_reloader) 

108 

109 

110@app.command() 

111def evaluate( 

112 predicted:Path=typer.Argument(..., help="The path to the TEI XML document with predictions from Rdgai to evaluate."), 

113 ground_truth:Path=typer.Argument(..., help="The path to the input TEI XML document to use as the ground truth for evaluation."), 

114 confusion_matrix:Path=typer.Option(None, help="Path to write the confusion matrix plot as a CSV file."), 

115 confusion_matrix_plot:Path=typer.Option(None, help="Path to write the confusion matrix plot as an HTML file."), 

116 report:Path=typer.Option(None, help="Path to write the report."), 

117): 

118 """ Evaluates the classifications in a predicted document against a ground truth document. """ 

119 predicted = Doc(predicted) 

120 ground_truth = Doc(ground_truth) 

121 

122 evaluate_docs(predicted, ground_truth, confusion_matrix=confusion_matrix, confusion_matrix_plot=confusion_matrix_plot, report=report) 

123 

124 

125@app.command() 

126def validate( 

127 ground_truth:Path=typer.Argument(..., help="The path to the input TEI XML document to use as the ground truth for evaluation."), 

128 output:Path=typer.Argument(..., help="The path to the output TEI XML file."), 

129 proportion:float=typer.Option(0.5, help="Proportion of classified pairs to use for validation."), 

130 api_key:str=typer.Option("", help="API key for the LLM."), 

131 llm:str=typer.Option(DEFAULT_MODEL_ID, help="ID of the language model to use."), 

132 temperature:float=typer.Option(0.1, help="Temperature for sampling from the language model."), 

133 examples:int=typer.Option(10, help="Number of examples to include in the prompt."), 

134 confusion_matrix:Path=typer.Option(None, help="Path to write the confusion matrix plot as a CSV file."), 

135 confusion_matrix_plot:Path=typer.Option(None, help="Path to write the confusion matrix plot as an HTML file."), 

136 seed:int=typer.Option(42, help="Seed for random sampling of validation pairs."), 

137 report:Path=typer.Option(None, help="Path to write the report."), 

138): 

139 """ Takes a ground truth document, chooses a proportion of classified pairs to validate against and outputs a report. """ 

140 ground_truth = Doc(ground_truth) 

141 

142 validate_fn( 

143 ground_truth, 

144 output, 

145 llm=llm, 

146 api_key=api_key, 

147 examples=examples, 

148 seed=seed, 

149 temperature=temperature, 

150 proportion=proportion, 

151 confusion_matrix=confusion_matrix, 

152 confusion_matrix_plot=confusion_matrix_plot, 

153 report=report, 

154 ) 

155 

156 

157@app.command() 

158def clean( 

159 doc:Path=typer.Argument(..., help="The path to the TEI XML document to clean."), 

160 output:Path=typer.Argument(None, help="The path to the output TEI XML file."), 

161 inplace: bool = typer.Option(False, "--inplace", "-i", help="Overwrite the input file."), 

162): 

163 """ Cleans a TEI XML file for common errors. """ 

164 output = get_output_path(doc, output, inplace) 

165 doc = Doc(doc) 

166 doc.clean(output=output) 

167 

168 

169@app.command() 

170def export( 

171 doc:Path=typer.Argument(..., help="The path to the TEI XML document to export."), 

172 output:Path=typer.Argument(..., help="The path to the output Excel file."), 

173): 

174 """ Exports pairs of readings with classifications from a TEI document to an Excel spreadsheet. """ 

175 doc = Doc(doc) 

176 export_variants_to_excel(doc, output) 

177 

178 

179@app.command() 

180def import_classifications( 

181 doc:Path=typer.Argument(..., help="The path to the base TEI XML document to use for importing the classifications from Excel."), 

182 spreadsheet:Path=typer.Argument(..., help="The path to the Excel file to import."), 

183 output:Path=typer.Argument(None, help="The path to the output TEI XML file."), 

184 inplace: bool = typer.Option(False, "--inplace", "-i", help="Overwrite the input file."), 

185 responsible:str=typer.Option("", help="The responsible party for the classifications. By default it is the name of the spreadsheet."), 

186): 

187 """ Imports classifications from a spreadsheet into a TEI document. """ 

188 doc = Doc(doc) 

189 output = get_output_path(doc, output, inplace) 

190 

191 if spreadsheet.suffix == ".xlsx": 

192 variants_df = pd.read_excel(spreadsheet, sheet_name="Variants", keep_default_na=False) 

193 elif spreadsheet.suffix == ".csv": 

194 variants_df = pd.read_csv(spreadsheet, keep_default_na=False) 

195 

196 # TODO add responsible to TEI header 

197 responsible = responsible or spreadsheet.stem 

198 responsible = responsible.replace(" ", "_") 

199 if not responsible.startswith("#"): 

200 responsible = "#" + responsible 

201 

202 import_classifications_from_dataframe(doc, variants_df, output, responsible=responsible) 

203 

204 

205@app.command() 

206def prompt_preamble( 

207 doc:Path=typer.Argument(..., help="The path to the TEI XML document to classify."), 

208 examples:int=typer.Option(10, help="Number of examples to include in the prompt."), 

209): 

210 """ Prints the prompt preamble for a TEI document for a given number of examples. """ 

211 doc = Doc(doc) 

212 template = build_preamble(doc, examples) 

213 print(template)