275 lines
11 KiB
Python
275 lines
11 KiB
Python
import os
|
|
import sys
|
|
import subprocess
|
|
import requests
|
|
import time
|
|
|
|
def check_cuda_11_8_setup():
|
|
"""Check if CUDA 11.8 is properly installed and configured"""
|
|
print("=== CHECKING CUDA 11.8 SETUP ===")
|
|
|
|
cuda_11_8_path = r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8'
|
|
|
|
# Check if CUDA 11.8 directory exists
|
|
if not os.path.exists(cuda_11_8_path):
|
|
print(f"✗ CUDA 11.8 not found at: {cuda_11_8_path}")
|
|
print("\nPlease install CUDA 11.8 from:")
|
|
print("https://developer.nvidia.com/cuda-11-8-0-download-archive")
|
|
return False
|
|
|
|
print(f"✓ CUDA 11.8 found at: {cuda_11_8_path}")
|
|
|
|
# Check for cuDNN
|
|
cudnn_dll = os.path.join(cuda_11_8_path, 'bin', 'cudnn64_8.dll')
|
|
if not os.path.exists(cudnn_dll):
|
|
print(f"✗ cuDNN 8.x not found at: {cudnn_dll}")
|
|
print("\nPlease install cuDNN 8.x for CUDA 11.8 from:")
|
|
print("https://developer.nvidia.com/rdp/cudnn-archive")
|
|
print("Then copy the DLLs to CUDA 11.8 bin directory")
|
|
return False
|
|
|
|
print(f"✓ cuDNN 8.x found at: {cudnn_dll}")
|
|
|
|
# Set environment for CUDA 11.8
|
|
os.environ['CUDA_PATH'] = cuda_11_8_path
|
|
os.environ['CUDA_HOME'] = cuda_11_8_path
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
|
|
os.environ['LIGHTRAG_OCR_ENGINE'] = 'paddleocr'
|
|
|
|
# Add CUDA 11.8 to PATH
|
|
cuda_bin = os.path.join(cuda_11_8_path, 'bin')
|
|
current_path = os.environ.get('PATH', '')
|
|
if cuda_bin not in current_path:
|
|
os.environ['PATH'] = cuda_bin + ';' + current_path
|
|
|
|
print("✓ Environment configured for CUDA 11.8")
|
|
return True
|
|
|
|
def test_paddlepaddle_gpu():
|
|
"""Test if PaddlePaddle can use GPU with CUDA 11.8"""
|
|
print("\n=== TESTING PADDLEPADDLE GPU ===")
|
|
|
|
try:
|
|
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 a simple GPU operation
|
|
x = paddle.to_tensor([1.0, 2.0, 3.0])
|
|
print(f"✓ GPU tensor test: {x.numpy()}")
|
|
return True
|
|
else:
|
|
print("✗ PaddlePaddle not compiled with CUDA")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"✗ PaddlePaddle test failed: {e}")
|
|
return False
|
|
|
|
def start_server_with_cuda11_8():
|
|
"""Start server with CUDA 11.8 environment"""
|
|
print("\n=== STARTING SERVER WITH CUDA 11.8 ===")
|
|
|
|
try:
|
|
# Use the zrun_cuda11.8.bat approach
|
|
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, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
|
|
# 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")
|
|
return None
|
|
|
|
except Exception as e:
|
|
print(f"✗ Failed to start server: {e}")
|
|
return None
|
|
|
|
def test_ocr_pdf_indexing():
|
|
"""Test OCR PDF indexing with CUDA 11.8"""
|
|
print("\n=== TESTING OCR PDF INDEXING WITH CUDA 11.8 ===")
|
|
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"\nUploading OCR PDF: ocr.pdf ({os.path.getsize('ocr.pdf')} bytes)")
|
|
with open('ocr.pdf', 'rb') as f:
|
|
files = {'file': ('ocr.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 processing
|
|
print(f"\n=== MONITORING OCR PROCESSING (CUDA 11.8) ===")
|
|
print("OCR processing with GPU acceleration...")
|
|
|
|
max_wait = 300 # 5 minutes for GPU processing
|
|
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
|
|
for doc in completed:
|
|
if doc.get('file_path') == 'ocr.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
|
|
print("\n=== TESTING SEARCH ===")
|
|
search_queries = [
|
|
"table", "data", "information", "content",
|
|
"scanned", "document", "text", "analysis"
|
|
]
|
|
|
|
total_results = 0
|
|
for query in search_queries:
|
|
print(f"Searching: '{query}'")
|
|
search_data = {'query': query, 'top_k': 3}
|
|
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)
|
|
print(f" Found {len(results)} results")
|
|
if results:
|
|
print(f" Best score: {results[0].get('score'):.3f}")
|
|
else:
|
|
print(f" Search failed: {search_response.text}")
|
|
except Exception as e:
|
|
print(f" Search error: {e}")
|
|
|
|
print(f"\n=== SUMMARY ===")
|
|
print(f"Total search results: {total_results}")
|
|
print("OCR PDF with scanned table successfully indexed and searchable!")
|
|
return True
|
|
|
|
# Check if failed
|
|
for doc in failed:
|
|
if doc.get('file_path') == 'ocr.pdf':
|
|
print(f"✗ OCR processing failed: {doc.get('error_msg', 'Unknown error')}")
|
|
return False
|
|
|
|
# Still processing
|
|
if elapsed % 30 == 0:
|
|
print(f" Processing... ({elapsed}s elapsed)")
|
|
|
|
time.sleep(5)
|
|
|
|
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 test: {e}")
|
|
return False
|
|
|
|
def main():
|
|
"""Main function to test CUDA 11.8 OCR setup"""
|
|
print("CUDA 11.8 OCR TEST SCRIPT")
|
|
print("=" * 50)
|
|
|
|
# Step 1: Check CUDA 11.8 setup
|
|
if not check_cuda_11_8_setup():
|
|
print("\n❌ CUDA 11.8 setup incomplete")
|
|
return
|
|
|
|
# Step 2: Test PaddlePaddle GPU
|
|
if not test_paddlepaddle_gpu():
|
|
print("\n❌ PaddlePaddle GPU test failed")
|
|
return
|
|
|
|
# Step 3: Start server with CUDA 11.8
|
|
server_process = start_server_with_cuda11_8()
|
|
if not server_process:
|
|
print("\n❌ Failed to start server")
|
|
return
|
|
|
|
try:
|
|
# Step 4: Test OCR PDF indexing
|
|
success = test_ocr_pdf_indexing()
|
|
|
|
if success:
|
|
print("\n🎉 SUCCESS: OCR PDF indexing with CUDA 11.8 completed!")
|
|
print(" GPU-accelerated OCR is working correctly.")
|
|
print(" The scanned table document has been processed and is searchable.")
|
|
else:
|
|
print("\n❌ OCR PDF indexing failed")
|
|
|
|
finally:
|
|
# Clean up
|
|
print("\nStopping server...")
|
|
server_process.terminate()
|
|
try:
|
|
server_process.wait(timeout=10)
|
|
except:
|
|
server_process.kill()
|
|
|
|
if __name__ == "__main__":
|
|
main() |