Update app.py

app.py

@@ -1,8 +1,10 @@
 import os
 from transformers import pipeline
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from queue import Queue
 import openai
+import gc
+import psutil
 
 # تعریف مدل‌ها
 MODEL_CONFIG = {
@@ -23,34 +25,54 @@ openai.api_key = os.getenv("OPENAI_API_KEY")
 if not openai.api_key:
     raise ValueError("OpenAI API key is missing! Set it as 'OPENAI_API_KEY' in environment variables.")
 
+# محدودیت حافظه (برحسب گیگابایت)
+MEMORY_LIMIT_GB = 15
+
+class MemoryManager:
+    @staticmethod
+    def check_memory_usage():
+        """بررسی مصرف حافظه و خروج در صورت تجاوز از محدودیت."""
+        process = psutil.Process(os.getpid())
+        memory_usage = process.memory_info().rss / (1024**3)  # تبدیل به گیگابایت
+        if memory_usage > MEMORY_LIMIT_GB:
+            raise MemoryError(f"Memory usage exceeded {MEMORY_LIMIT_GB} GB! Current usage: {memory_usage:.2f} GB.")
+
 class MultiModelSystem:
     def __init__(self):
         """مدیریت مدل‌های مختلف و پردازش موازی وظایف."""
         self.models = {}
         self.queue = Queue()
-        self.executor = ThreadPoolExecutor(max_workers=5)
+        self.executor = ThreadPoolExecutor(max_workers=3)  # محدود کردن وظایف همزمان
         self.load_models()
 
     def load_models(self):
-        """بارگذاری مدل‌ها به صورت موازی."""
+        """بارگذاری مدل‌ها به صورت دینامیک."""
         try:
             for task, model_id in MODEL_CONFIG.items():
                 if model_id == "OpenAI":
-                    self.models[task] = self.load_openai_model()
+                    self.models[task] = None  # مدل OpenAI نیازی به بارگذاری ندارد
                 else:
-                    self.models[task] = pipeline(
-                        task=self.get_task_type(task),
-                        model=model_id,
-                        use_auth_token=TOKEN
-                    )
-                print(f"Model '{task}' loaded successfully.")
+                    self.models[task] = model_id  # ذخیره نام مدل، نه بارگذاری کامل
         except Exception as e:
             print(f"Error loading models: {e}")
             raise
 
-    def load_openai_model(self):
-        """مدل ریاضی OpenAI را بارگذاری کنید."""
-        return "OpenAI (Math)"
+    def load_model_for_task(self, task):
+        """بارگذاری مدل موردنیاز فقط هنگام درخواست."""
+        model_id = self.models.get(task)
+        if not model_id:
+            return None
+        try:
+            MemoryManager.check_memory_usage()
+            return pipeline(
+                task=self.get_task_type(task),
+                model=model_id,
+                use_auth_token=TOKEN,
+                device=0 if psutil.cpu_count(logical=False) > 1 else -1  # استفاده از GPU در صورت وجود
+            )
+        except Exception as e:
+            print(f"Error loading model for task '{task}': {e}")
+            return None
 
     @staticmethod
     def get_task_type(task):
@@ -74,17 +96,24 @@ class MultiModelSystem:
                 if task == "math":
                     result = self.process_math_task(kwargs.get("text"))
                 else:
-                    result = self.models[task](**kwargs)
+                    model = self.load_model_for_task(task)
+                    if model:
+                        result = model(**kwargs)
+                    else:
+                        result = None
                 self.queue.put(result)
             except Exception as e:
                 print(f"Error processing task '{task}': {e}")
                 self.queue.put(None)
+            finally:
+                gc.collect()  # آزادسازی حافظه پس از هر وظیفه
 
         self.executor.submit(task_handler)
 
     def process_math_task(self, text):
         """پردازش ریاضی با OpenAI."""
         try:
+            MemoryManager.check_memory_usage()
             response = openai.Completion.create(
                 engine="text-davinci-003",
                 prompt=text,