printing extracted text12

main.py

@@ -3,21 +3,59 @@ import tempfile
 import gc
 import psutil
 import time
+import logging
+import queue
+
+# Set up logging first
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler()]
+)
+logger = logging.getLogger(__name__)
+
+# Create notification queue for real-time updates
+notification_queue = queue.Queue()
+
+def log_print(message, level="INFO"):
+    """Unified logging function"""
+    if level == "ERROR":
+        logger.error(message)
+    elif level == "WARNING":
+        logger.warning(message)
+    else:
+        logger.info(message)
+    # Also put the message in notification queue for frontend
+    notification_queue.put({
+        "type": level.lower(),
+        "message": message
+    })
 
 # Set environment variables before any other imports
 os.environ['TRANSFORMERS_CACHE'] = os.path.join(tempfile.gettempdir(), 'huggingface_cache')
 os.environ['HF_HOME'] = os.path.join(tempfile.gettempdir(), 'huggingface')
 os.environ['TORCH_HOME'] = os.path.join(tempfile.gettempdir(), 'torch')
 os.environ['XDG_CACHE_HOME'] = os.path.join(tempfile.gettempdir(), 'cache')
-
-# Create cache directories
-for cache_dir in [os.environ['TRANSFORMERS_CACHE'], os.environ['HF_HOME'], 
-                 os.environ['TORCH_HOME'], os.environ['XDG_CACHE_HOME']]:
-    os.makedirs(cache_dir, exist_ok=True)
+os.environ['SENTENCE_TRANSFORMERS_HOME'] = os.path.join(tempfile.gettempdir(), 'sentence_transformers')
+
+# Create all necessary cache directories
+cache_dirs = {
+    'transformers': os.environ['TRANSFORMERS_CACHE'],
+    'hf': os.environ['HF_HOME'],
+    'torch': os.environ['TORCH_HOME'],
+    'cache': os.environ['XDG_CACHE_HOME'],
+    'sentence_transformers': os.environ['SENTENCE_TRANSFORMERS_HOME']
+}
+
+for cache_name, cache_dir in cache_dirs.items():
+    try:
+        os.makedirs(cache_dir, exist_ok=True)
+        log_print(f"Created cache directory for {cache_name}: {cache_dir}")
+    except Exception as e:
+        log_print(f"Error creating {cache_name} cache directory: {e}", "ERROR")
 
 # Now import the rest of the dependencies
 import sys
-import logging
 from pathlib import Path
 from flask import Flask, request, jsonify, render_template, send_file, Response
 from werkzeug.utils import secure_filename
@@ -28,7 +66,6 @@ import io
 import base64
 from datetime import datetime
 import json
-import queue
 import threading
 from threading import Thread, Event
 import warnings
@@ -36,6 +73,14 @@ from flask_cors import CORS
 from dotenv import load_dotenv
 warnings.filterwarnings('ignore')
 
+# Import ML libraries
+import torch
+import nltk
+import gensim
+from gensim.models import FastText
+from sentence_transformers import SentenceTransformer
+from transformers import pipeline
+
 # Import ML libraries with timeout protection
 def import_with_timeout(import_statement, timeout=30):
     """Import a module with a timeout to prevent hanging"""
@@ -125,14 +170,6 @@ def ensure_directory(path):
             log_print(f"Error creating directory {path}: {nested_e}", "ERROR")
             raise
 
-# Simplified logging setup
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s',
-    handlers=[logging.StreamHandler(sys.stdout)]
-)
-logger = logging.getLogger(__name__)
-
 def get_or_load_model(model_name):
     """Get a model from cache or load it if not present"""
     if model_name not in global_models:
@@ -144,16 +181,84 @@ def get_or_load_model(model_name):
                 if not os.path.exists(model_path):
                     from gensim.downloader import load
                     log_print("Downloading fasttext model...")
-                    global_models[model_name] = load('fasttext-wiki-news-subwords-300')
+                    model = load('fasttext-wiki-news-subwords-300')
+                    # Move model to CPU explicitly
+                    if hasattr(model, 'to'):
+                        model = model.to('cpu')
+                    global_models[model_name] = model
                 else:
-                    global_models[model_name] = KeyedVectors.load_word2vec_format(model_path)
+                    model = KeyedVectors.load_word2vec_format(model_path)
+                    # Move model to CPU explicitly
+                    if hasattr(model, 'to'):
+                        model = model.to('cpu')
+                    global_models[model_name] = model
                 log_print(f"Successfully loaded {model_name} model")
+            elif model_name == 'vit':
+                try:
+                    from transformers import ViTImageProcessor, ViTModel
+                    log_print("Loading ViT model...")
+                    
+                    # Use a more common ViT model that's guaranteed to exist
+                    model_name = "google/vit-base-patch16-224-in21k"
+                    cache_dir = os.path.join(os.environ['TRANSFORMERS_CACHE'], 'vit-models')
+                    os.makedirs(cache_dir, exist_ok=True)
+                    
+                    try:
+                        # Try to load the processor first
+                        log_print("Loading ViT image processor...")
+                        processor = ViTImageProcessor.from_pretrained(model_name, 
+                                                                    cache_dir=cache_dir,
+                                                                    local_files_only=True)
+                        log_print("Loading ViT model from cache...")
+                        model = ViTModel.from_pretrained(model_name, 
+                                                       cache_dir=cache_dir,
+                                                       local_files_only=True)
+                        
+                        # Move model to CPU explicitly
+                        model = model.to('cpu')
+                        
+                        global_models['vit_processor'] = processor
+                        global_models['vit_model'] = model
+                        log_print("Successfully loaded ViT model from cache")
+                        
+                    except Exception as cache_error:
+                        log_print(f"Cache load failed ({str(cache_error)}), downloading model...")
+                        processor = ViTImageProcessor.from_pretrained(model_name, 
+                                                                    cache_dir=cache_dir,
+                                                                    local_files_only=False)
+                        model = ViTModel.from_pretrained(model_name, 
+                                                       cache_dir=cache_dir,
+                                                       local_files_only=False)
+                        
+                        # Move model to CPU explicitly
+                        model = model.to('cpu')
+                        
+                        global_models['vit_processor'] = processor
+                        global_models['vit_model'] = model
+                        log_print("Successfully downloaded and loaded ViT model")
+                    
+                except Exception as e:
+                    log_print(f"Error loading ViT model: {str(e)}", "ERROR")
+                    try:
+                        log_print("Trying alternative ViT model...")
+                        model_name = "google/vit-base-patch16-224"
+                        processor = ViTImageProcessor.from_pretrained(model_name)
+                        model = ViTModel.from_pretrained(model_name)
+                        
+                        # Move model to CPU explicitly
+                        model = model.to('cpu')
+                        
+                        global_models['vit_processor'] = processor
+                        global_models['vit_model'] = model
+                        log_print("Successfully loaded alternative ViT model")
+                    except Exception as alt_error:
+                        log_print(f"Error loading alternative ViT model: {str(alt_error)}", "ERROR")
+                        return None
             elif model_name == 'llm':
-                # Implement LLM model loading here
                 log_print("LLM model loading not implemented", "WARNING")
                 return None
         except Exception as e:
-            log_print(f"Error loading {model_name} model: {e}", "ERROR")
+            log_print(f"Error loading {model_name} model: {str(e)}", "ERROR")
             return None
     return global_models.get(model_name)
 
@@ -179,7 +284,11 @@ def initialize_resources():
 
         # Initialize models
         try:
+            # Load FastText first
             get_or_load_model('fasttext')
+            
+            # Then load ViT model
+            get_or_load_model('vit')
         except Exception as e:
             log_print(f"Warning: Could not preload models: {e}", "WARNING")