Update app.py

app.py

@@ -1,24 +1,20 @@
+import os
 import gc
 import psutil
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 
 class MultiModelSystem:
     """
-    سیستم چندمدلی برای مدیریت وظایف NLP با بهینه‌سازی مصرف حافظه.
+    سیستم چندمدلی با مدیریت حافظه و ذخیره‌سازی موقت مدل‌ها در دیسک.
     """
 
     def __init__(self, memory_limit_gb=15):
-        """
-        مقداردهی اولیه سیستم و تنظیم محدودیت حافظه.
-        :param memory_limit_gb: حداکثر میزان استفاده از حافظه.
-        """
         self.models = {}
         self.memory_limit_gb = memory_limit_gb
+        self.model_cache_dir = "model_cache"
+        os.makedirs(self.model_cache_dir, exist_ok=True)
 
     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)")
@@ -27,32 +23,27 @@ class MultiModelSystem:
 
     def load_model(self, task, model_id):
         """
-        بارگذاری مدل بر اساس وظیفه.
-        :param task: نوع وظیفه (مثلاً ترجمه).
-        :param model_id: شناسه مدل.
+        بارگذاری مدل از کش یا ذخیره‌سازی.
         """
+        cache_path = os.path.join(self.model_cache_dir, f"{task}.bin")
         if task not in self.models:
-            self.check_memory_usage()  # بررسی حافظه پیش از بارگذاری
-            print(f"Loading model for task '{task}' with ID '{model_id}'...")
-            if task == "translation":
+            self.check_memory_usage()
+            print(f"Loading model for task '{task}'...")
+            if os.path.exists(cache_path):
+                print(f"Loading model from cache: {cache_path}")
+                self.models[task] = joblib.load(cache_path)
+            else:
                 model = AutoModelForSeq2SeqLM.from_pretrained(
-                    model_id,
-                    torch_dtype="auto",  # بهینه‌سازی حافظه با FP16
-                    low_cpu_mem_usage=True
+                    model_id, torch_dtype="auto", 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)
+                joblib.dump(self.models[task], cache_path)
+                print(f"Model cached at {cache_path}")
 
     def unload_model(self, task):
         """
-        آزادسازی مدل برای مدیریت حافظه.
-        :param task: نوع وظیفه.
+        تخلیه مدل از حافظه.
         """
         if task in self.models:
             print(f"Unloading model for task '{task}'...")
@@ -61,45 +52,33 @@ class MultiModelSystem:
 
     def process_task(self, task, model_id, **kwargs):
         """
-        پردازش وظیفه با استفاده از مدل مناسب.
-        :param task: نوع وظیفه.
-        :param model_id: شناسه مدل.
-        :return: نتیجه پردازش.
+        پردازش وظیفه با بارگذاری موقت مدل.
         """
         self.load_model(task, model_id)
         model = self.models[task]
-
         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:
             raise ValueError(f"Unsupported task: {task}")
 
 if __name__ == "__main__":
-    # تنظیمات مدل‌ها
     MODEL_CONFIG = {
-        "translation": "PontifexMaximus/opus-mt-iir-en-finetuned-fa-to-en",
-        "qa": "HooshvareLab/bert-fa-base-uncased",
+        "translation": "Helsinki-NLP/opus-mt-en-ro",  # مدل سبک‌تر
+        "qa": "distilbert-base-uncased-distilled-squad",  # مدل فشرده
     }
 
-    # تعریف وظایف
     tasks = [
         {"task": "translation", "model_id": MODEL_CONFIG["translation"], "kwargs": {"text": "سلام دنیا!"}},
-        {"task": "qa", "model_id": MODEL_CONFIG["qa"], "kwargs": {"question": "پایتخت ایران کجاست؟", "context": "ایران کشوری در خاورمیانه است و پایتخت آن تهران است."}}
+        {"task": "qa", "model_id": MODEL_CONFIG["qa"], "kwargs": {"question": "What is AI?", "context": "AI is artificial intelligence."}}
     ]
 
-    # نمونه‌سازی سیستم
     system = MultiModelSystem(memory_limit_gb=15)
 
-    # پردازش وظایف
     for task_info in tasks:
         try:
             system.check_memory_usage()
@@ -108,4 +87,4 @@ if __name__ == "__main__":
         except Exception as e:
             print(f"Error during task '{task_info['task']}':", str(e))
         finally:
-            system.unload_model(task_info["task"])  # تخلیه مدل پس از اتمام
+            system.unload_model(task_info["task"])