apaxray commited on
Commit
da280fe
·
verified ·
1 Parent(s): 713c575

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +21 -42
app.py CHANGED
@@ -1,24 +1,20 @@
 
1
  import gc
2
  import psutil
3
  from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
4
 
5
  class MultiModelSystem:
6
  """
7
- سیستم چندمدلی برای مدیریت وظایف NLP با بهینه‌سازی مصرف حافظه.
8
  """
9
 
10
  def __init__(self, memory_limit_gb=15):
11
- """
12
- مقداردهی اولیه سیستم و تنظیم محدودیت حافظه.
13
- :param memory_limit_gb: حداکثر میزان استفاده از حافظه.
14
- """
15
  self.models = {}
16
  self.memory_limit_gb = memory_limit_gb
 
 
17
 
18
  def check_memory_usage(self):
19
- """
20
- بررسی میزان استفاده از حافظه.
21
- """
22
  mem = psutil.virtual_memory()
23
  used_gb = mem.used / (1024 ** 3)
24
  print(f"Memory usage: {mem.percent}% ({used_gb:.2f} GB used)")
@@ -27,32 +23,27 @@ class MultiModelSystem:
27
 
28
  def load_model(self, task, model_id):
29
  """
30
- بارگذاری مدل بر اساس وظیفه.
31
- :param task: نوع وظیفه (مثلاً ترجمه).
32
- :param model_id: شناسه مدل.
33
  """
 
34
  if task not in self.models:
35
- self.check_memory_usage() # بررسی حافظه پیش از بارگذاری
36
- print(f"Loading model for task '{task}' with ID '{model_id}'...")
37
- if task == "translation":
 
 
 
38
  model = AutoModelForSeq2SeqLM.from_pretrained(
39
- model_id,
40
- torch_dtype="auto", # بهینه‌سازی حافظه با FP16
41
- low_cpu_mem_usage=True
42
  )
43
  tokenizer = AutoTokenizer.from_pretrained(model_id)
44
  self.models[task] = pipeline("translation", model=model, tokenizer=tokenizer)
45
- elif task == "qa":
46
- model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
47
- tokenizer = AutoTokenizer.from_pretrained(model_id)
48
- self.models[task] = pipeline("question-answering", model=model, tokenizer=tokenizer)
49
- else:
50
- self.models[task] = pipeline(task, model=model_id)
51
 
52
  def unload_model(self, task):
53
  """
54
- آزادسازی مدل برای مدیریت حافظه.
55
- :param task: نوع وظیفه.
56
  """
57
  if task in self.models:
58
  print(f"Unloading model for task '{task}'...")
@@ -61,45 +52,33 @@ class MultiModelSystem:
61
 
62
  def process_task(self, task, model_id, **kwargs):
63
  """
64
- پردازش وظیفه با استفاده از مدل مناسب.
65
- :param task: نوع وظیفه.
66
- :param model_id: شناسه مدل.
67
- :return: نتیجه پردازش.
68
  """
69
  self.load_model(task, model_id)
70
  model = self.models[task]
71
-
72
  if task == "translation":
73
  text = kwargs.get("text", "")
74
- if not text:
75
- raise ValueError("No input text provided for translation task.")
76
  return model(text)
77
  elif task == "qa":
78
  question = kwargs.get("question", "")
79
  context = kwargs.get("context", "")
80
- if not question or not context:
81
- raise ValueError("Both 'question' and 'context' must be provided for QA task.")
82
  return model(question=question, context=context)
83
  else:
84
  raise ValueError(f"Unsupported task: {task}")
85
 
86
  if __name__ == "__main__":
87
- # تنظیمات مدل‌ها
88
  MODEL_CONFIG = {
89
- "translation": "PontifexMaximus/opus-mt-iir-en-finetuned-fa-to-en",
90
- "qa": "HooshvareLab/bert-fa-base-uncased",
91
  }
92
 
93
- # تعریف وظایف
94
  tasks = [
95
  {"task": "translation", "model_id": MODEL_CONFIG["translation"], "kwargs": {"text": "سلام دنیا!"}},
96
- {"task": "qa", "model_id": MODEL_CONFIG["qa"], "kwargs": {"question": "پایتخت ایران کجاست؟", "context": "ایران کشوری در خاورمیانه است و پایتخت آن تهران است."}}
97
  ]
98
 
99
- # نمونه‌سازی سیستم
100
  system = MultiModelSystem(memory_limit_gb=15)
101
 
102
- # پردازش وظایف
103
  for task_info in tasks:
104
  try:
105
  system.check_memory_usage()
@@ -108,4 +87,4 @@ if __name__ == "__main__":
108
  except Exception as e:
109
  print(f"Error during task '{task_info['task']}':", str(e))
110
  finally:
111
- system.unload_model(task_info["task"]) # تخلیه مدل پس از اتمام
 
1
+ import os
2
  import gc
3
  import psutil
4
  from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
5
 
6
  class MultiModelSystem:
7
  """
8
+ سیستم چندمدلی با مدیریت حافظه و ذخیره‌سازی موقت مدل‌ها در دیسک.
9
  """
10
 
11
  def __init__(self, memory_limit_gb=15):
 
 
 
 
12
  self.models = {}
13
  self.memory_limit_gb = memory_limit_gb
14
+ self.model_cache_dir = "model_cache"
15
+ os.makedirs(self.model_cache_dir, exist_ok=True)
16
 
17
  def check_memory_usage(self):
 
 
 
18
  mem = psutil.virtual_memory()
19
  used_gb = mem.used / (1024 ** 3)
20
  print(f"Memory usage: {mem.percent}% ({used_gb:.2f} GB used)")
 
23
 
24
  def load_model(self, task, model_id):
25
  """
26
+ بارگذاری مدل از کش یا ذخیره‌سازی.
 
 
27
  """
28
+ cache_path = os.path.join(self.model_cache_dir, f"{task}.bin")
29
  if task not in self.models:
30
+ self.check_memory_usage()
31
+ print(f"Loading model for task '{task}'...")
32
+ if os.path.exists(cache_path):
33
+ print(f"Loading model from cache: {cache_path}")
34
+ self.models[task] = joblib.load(cache_path)
35
+ else:
36
  model = AutoModelForSeq2SeqLM.from_pretrained(
37
+ model_id, torch_dtype="auto", low_cpu_mem_usage=True
 
 
38
  )
39
  tokenizer = AutoTokenizer.from_pretrained(model_id)
40
  self.models[task] = pipeline("translation", model=model, tokenizer=tokenizer)
41
+ joblib.dump(self.models[task], cache_path)
42
+ print(f"Model cached at {cache_path}")
 
 
 
 
43
 
44
  def unload_model(self, task):
45
  """
46
+ تخلیه مدل از حافظه.
 
47
  """
48
  if task in self.models:
49
  print(f"Unloading model for task '{task}'...")
 
52
 
53
  def process_task(self, task, model_id, **kwargs):
54
  """
55
+ پردازش وظیفه با بارگذاری موقت مدل.
 
 
 
56
  """
57
  self.load_model(task, model_id)
58
  model = self.models[task]
 
59
  if task == "translation":
60
  text = kwargs.get("text", "")
 
 
61
  return model(text)
62
  elif task == "qa":
63
  question = kwargs.get("question", "")
64
  context = kwargs.get("context", "")
 
 
65
  return model(question=question, context=context)
66
  else:
67
  raise ValueError(f"Unsupported task: {task}")
68
 
69
  if __name__ == "__main__":
 
70
  MODEL_CONFIG = {
71
+ "translation": "Helsinki-NLP/opus-mt-en-ro", # مدل سبک‌تر
72
+ "qa": "distilbert-base-uncased-distilled-squad", # مدل فشرده
73
  }
74
 
 
75
  tasks = [
76
  {"task": "translation", "model_id": MODEL_CONFIG["translation"], "kwargs": {"text": "سلام دنیا!"}},
77
+ {"task": "qa", "model_id": MODEL_CONFIG["qa"], "kwargs": {"question": "What is AI?", "context": "AI is artificial intelligence."}}
78
  ]
79
 
 
80
  system = MultiModelSystem(memory_limit_gb=15)
81
 
 
82
  for task_info in tasks:
83
  try:
84
  system.check_memory_usage()
 
87
  except Exception as e:
88
  print(f"Error during task '{task_info['task']}':", str(e))
89
  finally:
90
+ system.unload_model(task_info["task"])