Update app.py

app.py

@@ -20,7 +20,7 @@ from sklearn.metrics.pairwise import cosine_similarity
 from sklearn.feature_extraction.text import TfidfVectorizer
 import nltk
 from nltk.corpus import stopwords
-from huggingface_hub import download_url
+import wget
 
 nltk.download('stopwords')
 
@@ -250,7 +250,7 @@ async def load_models():
         model_path = os.path.join("models", model["filename"])
         if not os.path.exists(model_path):
             url = f"https://huggingface.co/{model['repo_id']}/resolve/main/{model['filename']}"
-            download_url(url, model_path)
+            wget.download(url, model_path)
         global_data['models'][model['name']] = Llama(model_path)
 
 async def generate_model_response(model, inputs):
@@ -297,6 +297,44 @@ iface = gr.Interface(
 def run_gradio():
     iface.launch(server_port=7862, prevent_thread_lock=True)
 
+def release_resources():
+    try:
+        torch.cuda.empty_cache()
+        gc.collect()
+    except Exception as e:
+        print(f"Failed to release resources: {e}")
+
+def resource_manager():
+    MAX_RAM_PERCENT = 1
+    MAX_CPU_PERCENT = 1
+    MAX_GPU_PERCENT = 1
+    MAX_RAM_MB = 1
+
+    while True:
+        try:
+            virtual_mem = psutil.virtual_memory()
+            current_ram_percent = virtual_mem.percent
+            current_ram_mb = virtual_mem.used / (1024 * 1024)
+
+            if current_ram_percent > MAX_RAM_PERCENT or current_ram_mb > MAX_RAM_MB:
+                release_resources()
+
+            current_cpu_percent = psutil.cpu_percent()
+            if current_cpu_percent > MAX_CPU_PERCENT:
+                psutil.Process(os.getpid()).nice()
+
+            if torch.cuda.is_available():
+                gpu = torch.cuda.current_device()
+                gpu_mem = torch.cuda.memory_percent(gpu)
+
+                if gpu_mem > MAX_GPU_PERCENT:
+                    release_resources()
+
+        except Exception as e:
+            print(f"Error in resource manager: {e}")
+
+resource_manager()
+
 if __name__ == "__main__":
     asyncio.run(load_models())
     Thread(target=run_uvicorn).start()