Update test.py

CosmosShadow · Jul 4, 2024 · e996b53 · e996b53
1 parent c392af9
commit e996b53
Showing 1 changed file with 10 additions and 9 deletions.
diff --git a/test/test.py b/test/test.py
@@ -4,10 +4,15 @@
 import dotenv
 dotenv.load_dotenv()
 
+pdf_path = '../examples/attention_is_all_you_need.pdf'
+output_dir = '../examples/attention_is_all_you_need/'   
+
+# 清空output_dir
+import shutil
+shutil.rmtree(output_dir, ignore_errors=True)
+
 def test_use_api_key():
     from gptpdf import parse_pdf
-    pdf_path = '../examples/attention_is_all_you_need.pdf'
-    output_dir = '../examples/attention_is_all_you_need/'
     api_key = os.getenv('OPENAI_API_KEY')
     base_url = os.getenv('OPENAI_API_BASE')
     # Manually provide OPENAI_API_KEY and OPEN_API_BASE
@@ -19,8 +24,7 @@ def test_use_api_key():
 
 def test_use_env():
     from gptpdf import parse_pdf
-    pdf_path = '../examples/attention_is_all_you_need.pdf'
-    output_dir = '../examples/attention_is_all_you_need/'
+
     # Use OPENAI_API_KEY and OPENAI_API_BASE from environment variables
     content, image_paths = parse_pdf(pdf_path, output_dir=output_dir, model='gpt-4o', verbose=True)
     print(content)
@@ -32,9 +36,6 @@ def test_azure():
     api_key = '8ef0b4df45e444079cd5a4xxxxx' # Azure API Key
     base_url = 'https://xxx.openai.azure.com/' # Azure API Base URL
     model = 'azure_xxxx' # azure_ with deploy ID name (not open ai model name), e.g. azure_cpgpt4
-
-    pdf_path = '../examples/attention_is_all_you_need.pdf'
-    output_dir = '../examples/attention_is_all_you_need/'
     # Use OPENAI_API_KEY and OPENAI_API_BASE from environment variables
     content, image_paths = parse_pdf(pdf_path, output_dir=output_dir, api_key=api_key, base_url=base_url, model=model, verbose=True)
     print(content)
@@ -44,6 +45,6 @@ def test_azure():
 
 
 if __name__ == '__main__':
-    # test_use_api_key()
+    test_use_api_key()
     # test_use_env()
-    test_azure()
+    # test_azure()