Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """ | |
| Test script for RunPod serverless endpoint | |
| """ | |
| import base64 | |
| import requests | |
| import json | |
| import sys | |
| def test_runpod_api(pdf_path, runpod_endpoint, runpod_api_key): | |
| """Test PDF conversion on RunPod""" | |
| # Read PDF and encode to base64 | |
| with open(pdf_path, 'rb') as f: | |
| pdf_data = f.read() | |
| pdf_base64 = base64.b64encode(pdf_data).decode('utf-8') | |
| # Prepare request | |
| headers = { | |
| 'Authorization': f'Bearer {runpod_api_key}', | |
| 'Content-Type': 'application/json' | |
| } | |
| payload = { | |
| 'input': { | |
| 'pdf_base64': pdf_base64, | |
| 'filename': pdf_path.split('/')[-1] | |
| } | |
| } | |
| print(f"Sending PDF to RunPod: {pdf_path}") | |
| print(f"File size: {len(pdf_data)} bytes") | |
| # Send request | |
| response = requests.post( | |
| f"{runpod_endpoint}/run", | |
| headers=headers, | |
| json=payload | |
| ) | |
| if response.status_code == 200: | |
| result = response.json() | |
| job_id = result.get('id') | |
| print(f"Job submitted: {job_id}") | |
| # Poll for result | |
| status_url = f"{runpod_endpoint}/status/{job_id}" | |
| while True: | |
| status_response = requests.get(status_url, headers=headers) | |
| if status_response.status_code == 200: | |
| status_data = status_response.json() | |
| if status_data['status'] == 'COMPLETED': | |
| output = status_data.get('output', {}) | |
| if output.get('status') == 'success': | |
| markdown = output.get('markdown', '') | |
| print(f"\nConversion successful!") | |
| print(f"Markdown length: {len(markdown)} characters") | |
| print(f"Pages: {output.get('pages', 'unknown')}") | |
| # Save result | |
| output_file = pdf_path.replace('.pdf', '_runpod.md') | |
| with open(output_file, 'w') as f: | |
| f.write(markdown) | |
| print(f"Saved to: {output_file}") | |
| else: | |
| print(f"Conversion failed: {output.get('error')}") | |
| break | |
| elif status_data['status'] == 'FAILED': | |
| print(f"Job failed: {status_data.get('error')}") | |
| break | |
| else: | |
| print(f"Status: {status_data['status']}") | |
| import time | |
| time.sleep(2) | |
| else: | |
| print(f"Status check failed: {status_response.status_code}") | |
| break | |
| else: | |
| print(f"Request failed: {response.status_code}") | |
| print(response.text) | |
| if __name__ == "__main__": | |
| # Example usage | |
| if len(sys.argv) < 3: | |
| print("Usage: python test_runpod.py <pdf_file> <runpod_endpoint> <api_key>") | |
| print("Example: python test_runpod.py test.pdf https://api.runpod.ai/v2/your-endpoint your-api-key") | |
| sys.exit(1) | |
| pdf_file = sys.argv[1] | |
| endpoint = sys.argv[2] | |
| api_key = sys.argv[3] if len(sys.argv) > 3 else input("RunPod API Key: ") | |
| test_runpod_api(pdf_file, endpoint, api_key) |