Update app.py

app.py

@@ -1,96 +1,111 @@
-from transformers import pipeline
-import os
-import openai
-import gc  # برای آزاد کردن حافظه
-import psutil  # برای مانیتور کردن حافظه
+import gc
+import psutil
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 
 class MultiModelSystem:
-    def __init__(self):
+    """
+    سیستم چندمدلی برای مدیریت وظایف NLP با بهینه‌سازی مصرف حافظه.
+    """
+
+    def __init__(self, memory_limit_gb=15):
+        """
+        مقداردهی اولیه سیستم و تنظیم محدودیت حافظه.
+        :param memory_limit_gb: حداکثر میزان استفاده از حافظه.
+        """
         self.models = {}
+        self.memory_limit_gb = memory_limit_gb
 
-    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=use_auth_token
-            )
+    def check_memory_usage(self):
+        """
+        بررسی میزان استفاده از حافظه.
+        """
+        mem = psutil.virtual_memory()
+        used_gb = mem.used / (1024 ** 3)
+        print(f"Memory usage: {mem.percent}% ({used_gb:.2f} GB used)")
+        if used_gb > self.memory_limit_gb:
+            raise MemoryError(f"Memory limit exceeded: {used_gb:.2f} GB used (limit: {self.memory_limit_gb} GB)")
+
+    def load_model(self, task, model_id):
+        """
+        بارگذاری مدل بر اساس وظیفه.
+        :param task: نوع وظیفه (مثلاً ترجمه).
+        :param model_id: شناسه مدل.
+        """
+        if task not in self.models:
+            self.check_memory_usage()  # بررسی حافظه پیش از بارگذاری
+            print(f"Loading model for task '{task}' with ID '{model_id}'...")
+            if task == "translation":
+                model = AutoModelForSeq2SeqLM.from_pretrained(
+                    model_id,
+                    torch_dtype="auto",  # بهینه‌سازی حافظه با FP16
+                    low_cpu_mem_usage=True
+                )
+                tokenizer = AutoTokenizer.from_pretrained(model_id)
+                self.models[task] = pipeline("translation", model=model, tokenizer=tokenizer)
+            elif task == "qa":
+                model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+                tokenizer = AutoTokenizer.from_pretrained(model_id)
+                self.models[task] = pipeline("question-answering", model=model, tokenizer=tokenizer)
+            else:
+                self.models[task] = pipeline(task, model=model_id)
 
     def unload_model(self, task):
-        """مدل بارگذاری شده را از حافظه پاک می‌کند."""
+        """
+        آزادسازی مدل برای مدیریت حافظه.
+        :param task: نوع وظیفه.
+        """
         if task in self.models:
+            print(f"Unloading model for task '{task}'...")
             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",
-        }
-        return task_map.get(task, "text-generation")
+            gc.collect()
 
     def process_task(self, task, model_id, **kwargs):
-        """مدیریت وظایف."""
-        if task not in self.models:
-            self.models[task] = self.load_model(task, model_id, use_auth_token=os.getenv("Passsssssss"))
+        """
+        پردازش وظیفه با استفاده از مدل مناسب.
+        :param task: نوع وظیفه.
+        :param model_id: شناسه مدل.
+        :return: نتیجه پردازش.
+        """
+        self.load_model(task, model_id)
         model = self.models[task]
 
-        # پردازش وظیفه
-        if task == "math":
-            return self.process_math_task(kwargs.get("text"))
+        if task == "translation":
+            text = kwargs.get("text", "")
+            if not text:
+                raise ValueError("No input text provided for translation task.")
+            return model(text)
+        elif task == "qa":
+            question = kwargs.get("question", "")
+            context = kwargs.get("context", "")
+            if not question or not context:
+                raise ValueError("Both 'question' and 'context' must be provided for QA task.")
+            return model(question=question, context=context)
         else:
-            return model(**kwargs)
+            raise ValueError(f"Unsupported task: {task}")
 
-    def process_math_task(self, text):
-        """مدیریت وظایف ریاضی OpenAI."""
-        try:
-            response = openai.ChatCompletion.create(
-                model="gpt-4",
-                messages=[{"role": "user", "content": text}]
-            )
-            return response['choices'][0]['message']['content'].strip()
-        except Exception as e:
-            print(f"Error processing math task: {e}")
-            return None
-
-# مانیتورینگ حافظه
-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", "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?"}},
+        {"task": "translation", "model_id": MODEL_CONFIG["translation"], "kwargs": {"text": "سلام دنیا!"}},
+        {"task": "qa", "model_id": MODEL_CONFIG["qa"], "kwargs": {"question": "پایتخت ایران کجاست؟", "context": "ایران کشوری در خاورمیانه است و پایتخت آن تهران است."}}
     ]
 
+    # نمونه‌سازی سیستم
+    system = MultiModelSystem(memory_limit_gb=15)
+
+    # پردازش وظایف
     for task_info in tasks:
-        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()  # دوباره بررسی حافظه
+        try:
+            system.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)
+        except Exception as e:
+            print(f"Error during task '{task_info['task']}':", str(e))
+        finally:
+            system.unload_model(task_info["task"])  # تخلیه مدل پس از اتمام