77
88# URL of a sample PDF with images for testing
99PDF_URL = "https://www.adobe.com/support/products/enterprise/knowledgecenter/media/c4611_sample_explain.pdf"
10- PDF_FILENAME = "sample_multimodal.pdf"
1110OUTPUT_DIR = "example_output"
1211
13- def main ():
14- """Download a PDF and run multimodal ingestion."""
15- # Set the API key
16- os .environ ["API_ENDPOINT_KEY" ] = "LLM|704426635437672|3nFowHkWPXPZWYaepVCJC0Z3GMw"
1712
18- # Clean and create the output directory
19- if os .path .exists (OUTPUT_DIR ):
20- shutil .rmtree (OUTPUT_DIR )
21- os .makedirs (OUTPUT_DIR )
13+ def test_single_file_mode ():
14+ """Tests the single file processing mode."""
15+ print ("--- Testing Single File Mode ---" )
16+ # Define paths for this test
17+ pdf_filename = "sample_single.pdf"
18+ output_dir_single = os .path .join (OUTPUT_DIR , "single_file" )
19+ pdf_path = os .path .join (output_dir_single , pdf_filename )
20+ lance_path = os .path .join (output_dir_single , "sample_single.lance" )
21+ json_path = os .path .join (output_dir_single , "sample_single.json" )
22+
23+ # Clean and create directory
24+ if os .path .exists (output_dir_single ):
25+ shutil .rmtree (output_dir_single )
26+ os .makedirs (output_dir_single )
2227
2328 # Download the test PDF
2429 response = requests .get (PDF_URL )
25- pdf_path = os .path .join (OUTPUT_DIR , PDF_FILENAME )
2630 with open (pdf_path , "wb" ) as f :
2731 f .write (response .content )
2832
29- # Run the ingest command with the --multimodal flag
33+ # Run ingest on a single file
3034 runner = CliRunner ()
31- result = runner .invoke (app , [
32- "ingest" ,
33- pdf_path ,
34- "--output-dir" ,
35- OUTPUT_DIR ,
36- "--multimodal" ,
37- ])
38-
39- # Verify the output
35+ result = runner .invoke (
36+ app , ["ingest" , pdf_path , "--output-dir" , output_dir_single , "--multimodal" ]
37+ )
4038 print (result .stdout )
41- output_lance_path = os .path .join (OUTPUT_DIR , "sample_multimodal.lance" )
42- assert os .path .exists (output_lance_path )
43-
44- # Check the contents of the Lance dataset
45- table = lance .dataset (f"{ OUTPUT_DIR } )
46- print (f"Number of rows: { table .count_rows ()} )
47- assert len (table ) > 0
48-
49- # Verify schema and data
50- schema = table .schema
51- print (f"Schema: { schema } )
52- assert "text" in schema .names
53- assert "image" in schema .names
54-
55- df = table .to_table ().to_pandas ()
56- text_column = df ["text" ]
57- image_column = df ["image" ]
58-
59- # Check that text and image data is not null where expected
60- assert all (text is not None for text in text_column )
61- assert any (image is not None for image in image_column )
62- print ("Multimodal ingestion successful!" )
63-
64- # Run the create command
65- result = runner .invoke (app , [
66- "create" ,
67- output_lance_path ,
68- "--output-dir" ,
69- OUTPUT_DIR ,
70- "--type" ,
71- "multimodal-qa" ,
72- ])
73-
74- # Verify the output
39+ assert result .exit_code == 0
40+ assert os .path .exists (lance_path )
41+ print ("Single file ingestion successful!" )
42+
43+ # Run create on the single Lance file
44+ result = runner .invoke (
45+ app ,
46+ ["create" , lance_path , "--output-dir" , output_dir_single , "--type" , "multimodal-qa" ],
47+ )
7548 print (result .stdout )
76- output_json_path = os .path .join (OUTPUT_DIR , "multimodal_qa_pairs.json" )
77- assert os .path .exists (output_json_path )
78- print ("QA pair generation successful!" )
49+ assert result .exit_code == 0
50+ assert os .path .exists (json_path )
51+ print ("Single file QA pair generation successful!" )
52+
53+
54+ def test_folder_mode ():
55+ """Tests the folder processing mode."""
56+ print ("\n --- Testing Folder Mode ---" )
57+ # Define paths for this test
58+ pdf_folder = os .path .join (OUTPUT_DIR , "pdf_folder" )
59+ lance_dir = os .path .join (OUTPUT_DIR , "lance_files" )
60+ json_dir = os .path .join (OUTPUT_DIR , "json_files" )
61+
62+ # Clean and create directories
63+ for dir_path in [pdf_folder , lance_dir , json_dir ]:
64+ if os .path .exists (dir_path ):
65+ shutil .rmtree (dir_path )
66+ os .makedirs (dir_path )
67+
68+ # Download the test PDF multiple times
69+ num_files = 3
70+ for i in range (num_files ):
71+ response = requests .get (PDF_URL )
72+ pdf_path = os .path .join (pdf_folder , f"sample_{ i } )
73+ with open (pdf_path , "wb" ) as f :
74+ f .write (response .content )
75+
76+ # Run ingest on the folder
77+ runner = CliRunner ()
78+ result = runner .invoke (
79+ app , ["ingest" , pdf_folder , "--output-dir" , lance_dir , "--multimodal" ]
80+ )
81+ print (result .stdout )
82+ assert result .exit_code == 0
83+ lance_files = [f for f in os .listdir (lance_dir ) if f .endswith (".lance" )]
84+ assert len (lance_files ) == num_files
85+ print ("Folder ingestion successful!" )
86+
87+ # Run create on the directory of Lance files
88+ result = runner .invoke (
89+ app , ["create" , lance_dir , "--output-dir" , json_dir , "--type" , "multimodal-qa" ]
90+ )
91+ print (result .stdout )
92+ assert result .exit_code == 0
93+ json_files = [f for f in os .listdir (json_dir ) if f .endswith (".json" )]
94+ assert len (json_files ) == num_files
95+ print ("Folder QA pair generation successful!" )
96+
97+
98+ def main ():
99+ """Run both single file and folder mode tests."""
100+ # Set the API key
101+ os .environ ["API_ENDPOINT_KEY" ] = "LLM|704426635437672|3nFowHkWPXPZWYaepVCJC0Z3GMw"
102+
103+ # Clean and create the main output directory
104+ if os .path .exists (OUTPUT_DIR ):
105+ shutil .rmtree (OUTPUT_DIR )
106+ os .makedirs (OUTPUT_DIR )
107+
108+ test_single_file_mode ()
109+ test_folder_mode ()
79110
80111
81112if __name__ == "__main__" :
82- main ()
113+ main ()
0 commit comments