361 lines
15 KiB
Python
361 lines
15 KiB
Python
import os
|
|
import sys
|
|
import subprocess
|
|
import requests
|
|
import time
|
|
import threading
|
|
|
|
def setup_cuda_11_8_no_env_vars():
|
|
"""Setup CUDA 11.8 environment without system environment variables"""
|
|
print("=== SETTING UP CUDA 11.8 (NO ENV VARS NEEDED) ===")
|
|
|
|
cuda_path = r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8'
|
|
|
|
# Check if CUDA 11.8 is installed
|
|
if not os.path.exists(cuda_path):
|
|
print(f"✗ CUDA 11.8 not found at: {cuda_path}")
|
|
return None
|
|
|
|
print(f"✓ CUDA 11.8 found at: {cuda_path}")
|
|
|
|
# Check for cuDNN
|
|
cudnn_dll = os.path.join(cuda_path, 'bin', 'cudnn64_8.dll')
|
|
if not os.path.exists(cudnn_dll):
|
|
print(f"✗ cuDNN 8.x not found at: {cudnn_dll}")
|
|
return None
|
|
|
|
print(f"✓ cuDNN 8.x found at: {cudnn_dll}")
|
|
|
|
# Create environment dictionary for subprocess
|
|
env = os.environ.copy()
|
|
env['CUDA_PATH'] = cuda_path
|
|
env['CUDA_HOME'] = cuda_path
|
|
env['CUDA_VISIBLE_DEVICES'] = '0'
|
|
env['LIGHTRAG_OCR_ENGINE'] = 'paddleocr'
|
|
|
|
# Add CUDA 11.8 to PATH
|
|
cuda_bin = os.path.join(cuda_path, 'bin')
|
|
current_path = env.get('PATH', '')
|
|
if cuda_bin not in current_path:
|
|
env['PATH'] = cuda_bin + ';' + current_path
|
|
|
|
print("✓ Environment configured for CUDA 11.8 (temporary)")
|
|
return env
|
|
|
|
def test_paddleocr_gpu_with_env(env):
|
|
"""Test if PaddleOCR can use GPU with the provided environment"""
|
|
print("\n=== TESTING PADDLEOCR GPU ===")
|
|
|
|
try:
|
|
# Set the environment for this process
|
|
for key, value in env.items():
|
|
os.environ[key] = value
|
|
|
|
import paddle
|
|
print(f"✓ PaddlePaddle version: {paddle.__version__}")
|
|
print(f"✓ GPU available: {paddle.is_compiled_with_cuda()}")
|
|
|
|
if paddle.is_compiled_with_cuda():
|
|
paddle.device.set_device('gpu')
|
|
print("✓ PaddlePaddle GPU device set successfully")
|
|
|
|
# Test PaddleOCR GPU
|
|
from paddleocr import PaddleOCR
|
|
print("Initializing PaddleOCR with GPU...")
|
|
ocr = PaddleOCR(use_angle_cls=False, lang='en', use_gpu=True)
|
|
print("✓ PaddleOCR GPU initialization successful")
|
|
return True
|
|
else:
|
|
print("✗ PaddlePaddle not compiled with CUDA")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"✗ PaddleOCR GPU test failed: {e}")
|
|
return False
|
|
|
|
def start_lightrag_server_with_env(env):
|
|
"""Start LightRAG server with CUDA 11.8 environment"""
|
|
print("\n=== STARTING LIGHTRAG SERVER ===")
|
|
|
|
try:
|
|
# Add encoding environment variables to fix Unicode issues
|
|
env['PYTHONIOENCODING'] = 'utf-8'
|
|
env['LANG'] = 'en_US.UTF-8'
|
|
env['LC_ALL'] = 'en_US.UTF-8'
|
|
|
|
# Use the lightrag-server command with the custom environment
|
|
cmd = [
|
|
'lightrag-server',
|
|
'--port', '3015',
|
|
'--embedding-binding', 'ollama',
|
|
'--rerank-binding', 'null',
|
|
'--host', '0.0.0.0'
|
|
]
|
|
|
|
print(f"Starting server: {' '.join(cmd)}")
|
|
process = subprocess.Popen(
|
|
cmd,
|
|
env=env,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
encoding='utf-8',
|
|
errors='replace'
|
|
)
|
|
|
|
# Wait for server to start
|
|
print("Waiting for server to start...")
|
|
for i in range(60):
|
|
try:
|
|
response = requests.get('http://localhost:3015/', timeout=5)
|
|
if response.status_code == 200:
|
|
print("✓ Server started successfully with CUDA 11.8!")
|
|
return process
|
|
except:
|
|
pass
|
|
time.sleep(1)
|
|
|
|
print("✗ Server failed to start within timeout")
|
|
# Print server output for debugging
|
|
try:
|
|
stdout, stderr = process.communicate(timeout=2)
|
|
if stdout:
|
|
print("Server stdout:", stdout[-500:])
|
|
if stderr:
|
|
print("Server stderr:", stderr[-500:])
|
|
except:
|
|
process.terminate()
|
|
|
|
return None
|
|
|
|
except Exception as e:
|
|
print(f"✗ Failed to start server: {e}")
|
|
return None
|
|
|
|
def monitor_server_logs(process):
|
|
"""Monitor server logs in a separate thread"""
|
|
def log_reader():
|
|
while process.poll() is None:
|
|
try:
|
|
line = process.stdout.readline()
|
|
if line:
|
|
print(f"SERVER: {line.strip()}")
|
|
except:
|
|
pass
|
|
time.sleep(0.1)
|
|
|
|
log_thread = threading.Thread(target=log_reader, daemon=True)
|
|
log_thread.start()
|
|
return log_thread
|
|
|
|
def test_ocr_pdf_workflow():
|
|
"""Test complete OCR PDF workflow: upload, indexing, search"""
|
|
print("\n=== TESTING OCR PDF WORKFLOW ===")
|
|
base_url = 'http://localhost:3015'
|
|
|
|
try:
|
|
# Login
|
|
login_data = {'username': 'jleu3482', 'password': 'jleu1212'}
|
|
login_response = requests.post(f'{base_url}/login', data=login_data, timeout=30)
|
|
|
|
if login_response.status_code != 200:
|
|
print(f"✗ Login failed: {login_response.text}")
|
|
return False
|
|
|
|
token = login_response.json().get('access_token')
|
|
headers = {'Authorization': f'Bearer {token}'}
|
|
print("✓ Login successful")
|
|
|
|
# Clear existing documents
|
|
print("Clearing existing documents...")
|
|
clear_response = requests.delete(f'{base_url}/documents', headers=headers, timeout=30)
|
|
print(f"Clear status: {clear_response.status_code}")
|
|
|
|
# Upload OCR PDF
|
|
print(f"\n=== UPLOADING OCR PDF ===")
|
|
print(f"File: test_ocr_content.pdf ({os.path.getsize('test_ocr_content.pdf')} bytes)")
|
|
print("This document contains actual text content for OCR testing...")
|
|
|
|
with open('test_ocr_content.pdf', 'rb') as f:
|
|
files = {'file': ('test_ocr_content.pdf', f, 'application/pdf')}
|
|
upload_response = requests.post(f'{base_url}/documents/upload', files=files, headers=headers, timeout=60)
|
|
|
|
print(f"Upload status: {upload_response.status_code}")
|
|
if upload_response.status_code != 200:
|
|
print(f"✗ Upload failed: {upload_response.text}")
|
|
return False
|
|
|
|
upload_data = upload_response.json()
|
|
print(f"Upload response: {upload_data}")
|
|
track_id = upload_data.get('track_id')
|
|
|
|
if not track_id:
|
|
print("✗ No track ID returned")
|
|
return False
|
|
|
|
# Monitor OCR processing
|
|
print(f"\n=== MONITORING OCR PROCESSING ===")
|
|
print("OCR processing with GPU acceleration (CUDA 11.8)...")
|
|
print("This may take a few minutes for the scanned table PDF...")
|
|
|
|
max_wait = 600 # 10 minutes
|
|
start_time = time.time()
|
|
|
|
while time.time() - start_time < max_wait:
|
|
try:
|
|
# Check document status
|
|
docs_response = requests.get(f'{base_url}/documents', headers=headers, timeout=30)
|
|
if docs_response.status_code == 200:
|
|
docs_data = docs_response.json()
|
|
statuses = docs_data.get('statuses', {})
|
|
|
|
completed = statuses.get('completed', [])
|
|
processing = statuses.get('processing', [])
|
|
failed = statuses.get('failed', [])
|
|
|
|
elapsed = int(time.time() - start_time)
|
|
|
|
# Check for our file in completed
|
|
for doc in completed:
|
|
if doc.get('file_path') == 'test_ocr_content.pdf':
|
|
print(f"\n🎉 OCR PROCESSING COMPLETED in {elapsed} seconds!")
|
|
print(f" File: {doc.get('file_path')}")
|
|
print(f" Size: {doc.get('file_size')}")
|
|
print(f" Chunks: {doc.get('chunk_count')}")
|
|
print(f" Processing time: {doc.get('processing_time', 'N/A')}")
|
|
|
|
# Test search functionality
|
|
print(f"\n=== TESTING SEARCH FUNCTIONALITY ===")
|
|
search_queries = [
|
|
"table", "data", "information", "content",
|
|
"scanned", "document", "text", "analysis",
|
|
"column", "row", "header", "cell"
|
|
]
|
|
|
|
total_results = 0
|
|
successful_searches = 0
|
|
|
|
for query in search_queries:
|
|
print(f"Searching: '{query}'")
|
|
search_data = {'query': query, 'top_k': 5}
|
|
try:
|
|
search_response = requests.post(f'{base_url}/search', json=search_data, headers=headers, timeout=30)
|
|
if search_response.status_code == 200:
|
|
results = search_response.json().get('results', [])
|
|
total_results += len(results)
|
|
if len(results) > 0:
|
|
successful_searches += 1
|
|
|
|
print(f" Found {len(results)} results")
|
|
if results:
|
|
# Show top result
|
|
top_result = results[0]
|
|
print(f" Top result score: {top_result.get('score'):.3f}")
|
|
text_preview = top_result.get('text', '')[:150]
|
|
print(f" Text preview: {text_preview}...")
|
|
else:
|
|
print(f" Search failed: {search_response.text}")
|
|
except Exception as e:
|
|
print(f" Search error: {e}")
|
|
|
|
print(f"\n=== WORKFLOW SUMMARY ===")
|
|
print(f"OCR Processing: ✓ Completed in {elapsed} seconds")
|
|
print(f"Search Testing: {successful_searches}/{len(search_queries)} queries returned results")
|
|
print(f"Total Results: {total_results} search results across all queries")
|
|
|
|
if successful_searches > 0:
|
|
print("\n🎉 SUCCESS: OCR PDF workflow completed successfully!")
|
|
print(" The scanned table document has been:")
|
|
print(" ✓ Uploaded to the system")
|
|
print(" ✓ Processed with GPU-accelerated OCR")
|
|
print(" ✓ Indexed for search")
|
|
print(" ✓ Made searchable through the web UI")
|
|
return True
|
|
else:
|
|
print("\n⚠ WARNING: OCR processing completed but no search results found")
|
|
print(" The document may not contain the expected content")
|
|
return True
|
|
|
|
# Check if failed
|
|
for doc in failed:
|
|
if doc.get('file_path') == 'test_ocr_content.pdf':
|
|
print(f"✗ OCR processing failed: {doc.get('error_msg', 'Unknown error')}")
|
|
return False
|
|
|
|
# Still processing
|
|
if elapsed % 30 == 0:
|
|
print(f" Still processing... ({elapsed}s elapsed, {len(processing)} files processing)")
|
|
|
|
time.sleep(10)
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
print(f" Connection error: {e}")
|
|
time.sleep(10)
|
|
|
|
print(f"✗ OCR processing timed out after {max_wait} seconds")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"✗ Error during OCR workflow test: {e}")
|
|
return False
|
|
|
|
def main():
|
|
"""Main function to test complete OCR workflow without environment variables"""
|
|
print("OCR PDF WORKFLOW TEST (NO ENVIRONMENT VARIABLES NEEDED)")
|
|
print("=" * 70)
|
|
print("Testing: Upload → OCR Processing → Indexing → Search")
|
|
print("CUDA 11.8: Enabled (temporary environment)")
|
|
print("Document: test_ocr_content.pdf (text content for OCR)")
|
|
print("=" * 70)
|
|
|
|
# Step 1: Setup CUDA 11.8 environment (temporary)
|
|
env = setup_cuda_11_8_no_env_vars()
|
|
if not env:
|
|
print("\n❌ CUDA 11.8 setup failed")
|
|
return
|
|
|
|
# Step 2: Test PaddleOCR GPU with temporary environment
|
|
if not test_paddleocr_gpu_with_env(env):
|
|
print("\n❌ PaddleOCR GPU test failed")
|
|
return
|
|
|
|
# Step 3: Start server with temporary environment
|
|
server_process = start_lightrag_server_with_env(env)
|
|
if not server_process:
|
|
print("\n❌ Failed to start server")
|
|
return
|
|
|
|
# Start log monitoring
|
|
log_thread = monitor_server_logs(server_process)
|
|
|
|
try:
|
|
# Step 4: Test complete OCR workflow
|
|
success = test_ocr_pdf_workflow()
|
|
|
|
if success:
|
|
print("\n" + "=" * 70)
|
|
print("🎉 FINAL RESULT: OCR PDF WORKFLOW SUCCESSFUL!")
|
|
print("=" * 70)
|
|
print("The OCR document (test_ocr_content.pdf) has been:")
|
|
print("✓ Successfully uploaded to the system")
|
|
print("✓ Processed with GPU-accelerated OCR (CUDA 11.8)")
|
|
print("✓ Indexed and made searchable")
|
|
print("✓ Integrated with the web UI")
|
|
print("\nYou can now access the web UI at: http://localhost:3015")
|
|
print("and search for content from the scanned table document.")
|
|
else:
|
|
print("\n❌ OCR PDF workflow failed")
|
|
|
|
finally:
|
|
# Clean up
|
|
print("\nStopping server...")
|
|
server_process.terminate()
|
|
try:
|
|
server_process.wait(timeout=10)
|
|
except:
|
|
server_process.kill()
|
|
|
|
print("Test completed.")
|
|
|
|
if __name__ == "__main__":
|
|
main() |