Update app.py

app.py

@@ -14,6 +14,9 @@ from nltk.corpus import stopwords
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 import nltk
+import uvicorn
+import psutil
+import torch
 
 nltk.download('punkt')
 nltk.download('stopwords')
@@ -52,7 +55,7 @@ class ModelManager:
             if model_name not in self.models:
                 try:
                     model_path = hf_hub_download(repo_id=config['repo_id'], use_auth_token=HUGGINGFACE_TOKEN)
-                    model = Llama.from_file(model_path)
+                    model = Llama.from_file(model_path, n_ctx=512, n_gpu=1)
                     self.models[model_name] = model
                 except Exception as e:
                     self.models[model_name] = None
@@ -82,7 +85,7 @@ async def process_message(message: str) -> dict:
     inputs = message.strip()
     responses = {}
 
-    with ThreadPoolExecutor(max_workers=len(global_data['model_configs'])) as executor:
+    with ThreadPoolExecutor(max_workers=min(len(global_data['model_configs']), 4)) as executor:
         futures = [executor.submit(generate_model_response, model_manager.get_model(config['name']), inputs) for config in global_data['model_configs'] if model_manager.get_model(config['name'])]
         for i, future in enumerate(tqdm(as_completed(futures), total=len(futures), desc="Generating responses")):
             try:
@@ -127,6 +130,46 @@ async def startup_event():
 async def shutdown_event():
     gc.collect()
 
+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}")
+
 if __name__ == "__main__":
+    import threading
+    resource_thread = threading.Thread(target=resource_manager)
+    resource_thread.daemon = True
+    resource_thread.start()
     port = int(os.environ.get("PORT", 7860))
     uvicorn.run(app, host="0.0.0.0", port=port)