Update app.py

app.py

@@ -9,7 +9,7 @@ from datetime import datetime
 from pathlib import Path
 import gc
 import torch
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 from transformers import AutoModel, AutoTokenizer
 from sentence_transformers import SentenceTransformer
 import numpy as np
@@ -78,7 +78,7 @@ def initialize_model():
 def handle_gpu_operation(func):
     try:
         start_time = datetime.now()
-        with autocast(enabled=torch.cuda.is_available()):
+        with autocast(device_type='cuda', dtype=torch.float16):
             result = func()
         end_time = datetime.now()
         duration = (end_time - start_time).total_seconds()
@@ -230,20 +230,18 @@ def browse_outputs():
         return "Error opening file browser"
 
 def download_results():
+    required_files = ["embeddings.npy", "chunks.txt"]
+    missing = [f for f in required_files if not os.path.exists(os.path.join(OUTPUTS_DIR, f))]
+    if missing:
+        logger.error(f"Missing files: {missing}")
+        return None
+    
     try:
-        output_files = [
-            os.path.join(OUTPUTS_DIR, "embeddings.npy"),
-            os.path.join(OUTPUTS_DIR, "chunks.txt")
-        ]
-        
-        # Create a temporary zip file
-        with tempfile.TemporaryDirectory() as temp_dir:
-            zip_path = os.path.join(temp_dir, "results.zip")
-            with zipfile.ZipFile(zip_path, 'w') as zipf:
-                for file in output_files:
-                    if os.path.exists(file):
-                        zipf.write(file, os.path.basename(file))
-            return zip_path
+        zip_path = os.path.join(OUTPUTS_DIR, "results.zip")
+        with zipfile.ZipFile(zip_path, 'w') as zipf:
+            for file in required_files:
+                zipf.write(os.path.join(OUTPUTS_DIR, file), file)
+        return zip_path
     except Exception as e:
         logger.error(f"Error creating download: {str(e)}")
         return None