Update app.py

app.py

@@ -1,150 +1,96 @@
-import os
 from transformers import pipeline
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from queue import Queue
+import os
 import openai
-import gc
-import psutil
-
-# تعریف مدل‌ها
-MODEL_CONFIG = {
-    "translation": "PontifexMaximus/opus-mt-iir-en-finetuned-fa-to-en",  # ترجمه
-    "qa": "HooshvareLab/bert-fa-base-uncased",  # پاسخ به سوال
-    "math": "OpenAI",  # ریاضی با OpenAI
-    "persian_nlp": "HooshvareLab/bert-fa-zwnj-base",  # پردازش زبان فارسی
-    "custom_ai": "universitytehran/PersianMind-v1.0",  # شخصی‌سازی شده
-}
-
-# بارگذاری توکن از متغیر محیطی
-TOKEN = os.getenv("Passsssssss")
-if not TOKEN:
-    raise ValueError("API token is missing! Set it as 'Passsssssss' in environment variables.")
-
-# تنظیم کلید API OpenAI
-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.")
+import gc  # برای آزاد کردن حافظه
+import psutil  # برای مانیتور کردن حافظه
 
 class MultiModelSystem:
     def __init__(self):
-        """مدیریت مدل‌های مختلف و پردازش موازی وظایف."""
         self.models = {}
-        self.queue = Queue()
-        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] = None  # مدل OpenAI نیازی به بارگذاری ندارد
-                else:
-                    self.models[task] = model_id  # ذخیره نام مدل، نه بارگذاری کامل
-        except Exception as e:
-            print(f"Error loading models: {e}")
-            raise
 
-    def load_model_for_task(self, task):
-        """بارگذاری مدل موردنیاز فقط هنگام درخواست."""
-        model_id = self.models.get(task)
-        if not model_id:
-            return None
-        try:
-            MemoryManager.check_memory_usage()
+    def load_model(self, task, model_id, use_auth_token=None):
+        """مدل‌ها را به صورت lazy بارگذاری می‌کند."""
+        if model_id == "OpenAI":
+            return self.load_openai_model()
+        else:
             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 در صورت وجود
+                use_auth_token=use_auth_token
             )
-        except Exception as e:
-            print(f"Error loading model for task '{task}': {e}")
-            return None
+
+    def unload_model(self, task):
+        """مدل بارگذاری شده را از حافظه پاک می‌کند."""
+        if task in self.models:
+            del self.models[task]
+            gc.collect()  # جمع‌آوری حافظه
+
+    @staticmethod
+    def load_openai_model():
+        """مدل ریاضی OpenAI."""
+        return "OpenAI (Math)"
 
     @staticmethod
     def get_task_type(task):
-        """بازگرداندن نوع وظیفه براساس مدل."""
         task_map = {
             "translation": "translation",
             "qa": "question-answering",
             "persian_nlp": "text-classification",
             "custom_ai": "text-generation",
-            "math": "text-generation"
+            "math": "text-generation",
         }
         return task_map.get(task, "text-generation")
 
-    def process_task(self, task, **kwargs):
-        """پردازش وظایف به صورت غیرهمزمان."""
+    def process_task(self, task, model_id, **kwargs):
+        """مدیریت وظایف."""
         if task not in self.models:
-            raise ValueError(f"Model for task '{task}' not loaded.")
-
-        def task_handler():
-            try:
-                if task == "math":
-                    result = self.process_math_task(kwargs.get("text"))
-                else:
-                    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.models[task] = self.load_model(task, model_id, use_auth_token=os.getenv("Passsssssss"))
+        model = self.models[task]
 
-        self.executor.submit(task_handler)
+        # پردازش وظیفه
+        if task == "math":
+            return self.process_math_task(kwargs.get("text"))
+        else:
+            return model(**kwargs)
 
     def process_math_task(self, text):
-        """پردازش ریاضی با OpenAI."""
+        """مدیریت وظایف ریاضی OpenAI."""
         try:
-            MemoryManager.check_memory_usage()
-            response = openai.Completion.create(
-                engine="text-davinci-003",
-                prompt=text,
-                max_tokens=100
+            response = openai.ChatCompletion.create(
+                model="gpt-4",
+                messages=[{"role": "user", "content": text}]
             )
-            return response.choices[0].text.strip()
+            return response['choices'][0]['message']['content'].strip()
         except Exception as e:
-            print(f"Error processing math task with OpenAI: {e}")
+            print(f"Error processing math task: {e}")
             return None
 
-    def get_result(self):
-        """دریافت نتیجه از صف."""
-        return self.queue.get()
+# مانیتورینگ حافظه
+def check_memory_usage():
+    mem = psutil.virtual_memory()
+    print(f"Memory usage: {mem.percent}% ({mem.used / (1024 ** 3):.2f} GB used)")
 
 # نمونه استفاده
 if __name__ == "__main__":
+    MODEL_CONFIG = {
+        "translation": "PontifexMaximus/opus-mt-iir-en-finetuned-fa-to-en",
+        "qa": "HooshvareLab/bert-fa-base-uncased",
+        "math": "OpenAI",
+        "persian_nlp": "HooshvareLab/bert-fa-zwnj-base",
+        "custom_ai": "universitytehran/PersianMind-v1.0",
+    }
+
     system = MultiModelSystem()
 
-    # وظایف مختلف
     tasks = [
-        {"task": "translation", "kwargs": {"text": "سلام دنیا!", "src_lang": "fa", "tgt_lang": "en"}},
-        {"task": "qa", "kwargs": {"question": "پایتخت ایران چیست؟", "context": "ایران کشوری در خاورمیانه است و پایتخت آن تهران است."}},
-        {"task": "math", "kwargs": {"text": "What is the integral of x^2?"}},
-        {"task": "persian_nlp", "kwargs": {"text": "این یک جمله فارسی است."}},
-        {"task": "custom_ai", "kwargs": {"text": "تحلیل متن‌های تاریخی طبری."}}
+        {"task": "translation", "model_id": MODEL_CONFIG["translation"], "kwargs": {"text": "سلام دنیا!", "src_lang": "fa", "tgt_lang": "en"}},
+        {"task": "qa", "model_id": MODEL_CONFIG["qa"], "kwargs": {"question": "پایتخت ایران چیست؟", "context": "ایران کشوری در خاورمیانه است و پایتخت آن تهران است."}},
+        {"task": "math", "model_id": MODEL_CONFIG["math"], "kwargs": {"text": "What is the integral of x^2?"}},
     ]
 
-    # ارسال وظایف برای پردازش
     for task_info in tasks:
-        system.process_task(task_info["task"], **task_info["kwargs"])
-
-    # بازیابی نتایج
-    for _ in range(len(tasks)):
-        print("Result:", system.get_result())
+        check_memory_usage()  # نمایش میزان حافظه
+        result = system.process_task(task_info["task"], task_info["model_id"], **task_info["kwargs"])
+        print(f"Result for task '{task_info['task']}':", result)
+        system.unload_model(task_info["task"])  # آزادسازی حافظه
+        check_memory_usage()  # دوباره بررسی حافظه