Update app.py

app.py

@@ -4,7 +4,7 @@ import logging
 import numpy as np
 from sentence_transformers import SentenceTransformer
 import torch
-from torch.amp import autocast  # Corrected import for autocast
+from torch.amp import autocast
 from spaces import GPU
 import json  # Import json for direct JSON output in UI
 
@@ -32,26 +32,34 @@ logger = logging.getLogger(__name__)
 
 # Model initialization
 model = None
+model_initialization_error = "" # Global variable to store initialization error
 
 def initialize_model():
-    global model
+    global model, model_initialization_error
     try:
         if model is None:
             model = SentenceTransformer(EMBEDDING_MODEL_NAME, cache_folder=CACHE_DIR, use_auth_token=HF_TOKEN)
             logger.info(f"Initialized model: {EMBEDDING_MODEL_NAME}")
-        return True
+            model_initialization_error = "" # Clear any previous error
+            return True, "" # Return success and no error message
+        return True, "" # Already initialized, return success and no error
     except Exception as e:
-        logger.error(f"Model initialization failed: {str(e)}")
-        return False
+        error_msg = f"Model initialization failed: {str(e)}"
+        logger.error(error_msg)
+        model_initialization_error = error_msg # Store error message
+        return False, error_msg # Return failure and error message
+
 
 @GPU()
 def generate_embedding(text, focus):
-    global model
+    global model, model_initialization_error
     if model is None:
-        initialize_model()
+        success, error_message = initialize_model() # Call initialize_model and get status
+        if not success:
+            return "", error_message # Return initialization error to UI
 
     try:
-        with torch.amp.autocast('cuda'): # Corrected autocast usage
+        with torch.amp.autocast('cuda'):
             embedding_vector = model.encode([text])[0].tolist() # Get embedding as list
         # Convert embedding to JSON string for direct display in UI
         embedding_json_str = json.dumps(embedding_vector)
@@ -87,9 +95,11 @@ def convert_to_json(embedding_json, name): # Expect JSON string as input
 
 @GPU()
 def process_files(files, focus):
-    global model
+    global model, model_initialization_error
     if model is None:
-        initialize_model()
+        success, error_message = initialize_model() # Call initialize_model and get status
+        if not success:
+            return "", error_message # Return initialization error to UI
 
     try:
         all_embeddings = []
@@ -98,7 +108,7 @@ def process_files(files, focus):
             try:
                 with open(file.name, 'r') as f:
                     text = f.read()
-                with torch.amp.autocast('cuda'): # Corrected autocast usage
+                with torch.amp.autocast('cuda'):
                     embedding = model.encode([text])[0].tolist()
                 all_embeddings.append(embedding)
                 file_statuses.append(f"File '{file.name}' processed successfully.")
@@ -123,6 +133,8 @@ def create_gradio_interface():
     with gr.Blocks() as demo:
         gr.Markdown("## Text Embedding Generator")
 
+        initialization_status_box = gr.Textbox(label="Initialization Status", value=model_initialization_error, visible=False) # Hidden box to hold init error
+
         with gr.Row():
             text_input = gr.Textbox(label="Enter Text")
             focus_input = gr.Textbox(label="Main Focus of Embedding (e.g., company structure, staff positions, etc.)")
@@ -149,6 +161,11 @@ def create_gradio_interface():
         process_output = gr.Textbox(label="Processed Files (Embeddings JSON - limited display)", lines=3) # Limited lines for process output
         process_status = gr.Textbox(label="File Processing Status") # Status for file processing
 
+        demo.load( # Call initialize_model on app load
+            lambda: ("", model_initialization_error), # Dummy output for other components, error for initialization_status_box
+            outputs=[status_box, initialization_status_box] # status_box for general messages, init status for hidden box
+        )
+
 
         generate_button.click(
             generate_embedding,
@@ -183,5 +200,10 @@ def create_gradio_interface():
     return demo
 
 if __name__ == "__main__":
+    # Explicitly initialize the model at app startup and check for errors
+    initialization_success, initialization_error_message = initialize_model()
+    if not initialization_success:
+        print(f"App startup failed due to model initialization error:\n{initialization_error_message}") # Print to console for startup errors
+
     demo = create_gradio_interface()
     demo.launch(server_name="0.0.0.0")