apaxray commited on
Commit
aaed90f
·
verified ·
1 Parent(s): cdfaa37

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +56 -110
app.py CHANGED
@@ -1,150 +1,96 @@
1
- import os
2
  from transformers import pipeline
3
- from concurrent.futures import ThreadPoolExecutor, as_completed
4
- from queue import Queue
5
  import openai
6
- import gc
7
- import psutil
8
-
9
- # تعریف مدل‌ها
10
- MODEL_CONFIG = {
11
- "translation": "PontifexMaximus/opus-mt-iir-en-finetuned-fa-to-en", # ترجمه
12
- "qa": "HooshvareLab/bert-fa-base-uncased", # پاسخ به سوال
13
- "math": "OpenAI", # ریاضی با OpenAI
14
- "persian_nlp": "HooshvareLab/bert-fa-zwnj-base", # پردازش زبان فارسی
15
- "custom_ai": "universitytehran/PersianMind-v1.0", # شخصی‌سازی شده
16
- }
17
-
18
- # بارگذاری توکن از متغیر محیطی
19
- TOKEN = os.getenv("Passsssssss")
20
- if not TOKEN:
21
- raise ValueError("API token is missing! Set it as 'Passsssssss' in environment variables.")
22
-
23
- # تنظیم کلید API OpenAI
24
- openai.api_key = os.getenv("OPENAI_API_KEY")
25
- if not openai.api_key:
26
- raise ValueError("OpenAI API key is missing! Set it as 'OPENAI_API_KEY' in environment variables.")
27
-
28
- # محدودیت حافظه (برحسب گیگابایت)
29
- MEMORY_LIMIT_GB = 15
30
-
31
- class MemoryManager:
32
- @staticmethod
33
- def check_memory_usage():
34
- """بررسی مصرف حافظه و خروج در صورت تجاوز از محدودیت."""
35
- process = psutil.Process(os.getpid())
36
- memory_usage = process.memory_info().rss / (1024**3) # تبدیل به گیگابایت
37
- if memory_usage > MEMORY_LIMIT_GB:
38
- raise MemoryError(f"Memory usage exceeded {MEMORY_LIMIT_GB} GB! Current usage: {memory_usage:.2f} GB.")
39
 
40
  class MultiModelSystem:
41
  def __init__(self):
42
- """مدیریت مدل‌های مختلف و پردازش موازی وظایف."""
43
  self.models = {}
44
- self.queue = Queue()
45
- self.executor = ThreadPoolExecutor(max_workers=3) # محدود کردن وظایف همزمان
46
- self.load_models()
47
-
48
- def load_models(self):
49
- """بارگذاری مدل‌ها به صورت دینامیک."""
50
- try:
51
- for task, model_id in MODEL_CONFIG.items():
52
- if model_id == "OpenAI":
53
- self.models[task] = None # مدل OpenAI نیازی به بارگذاری ندارد
54
- else:
55
- self.models[task] = model_id # ذخیره نام مدل، نه بارگذاری کامل
56
- except Exception as e:
57
- print(f"Error loading models: {e}")
58
- raise
59
 
60
- def load_model_for_task(self, task):
61
- """بارگذاری مدل موردنیاز فقط هنگام درخواست."""
62
- model_id = self.models.get(task)
63
- if not model_id:
64
- return None
65
- try:
66
- MemoryManager.check_memory_usage()
67
  return pipeline(
68
  task=self.get_task_type(task),
69
  model=model_id,
70
- use_auth_token=TOKEN,
71
- device=0 if psutil.cpu_count(logical=False) > 1 else -1 # استفاده از GPU در صورت وجود
72
  )
73
- except Exception as e:
74
- print(f"Error loading model for task '{task}': {e}")
75
- return None
 
 
 
 
 
 
 
 
76
 
77
  @staticmethod
78
  def get_task_type(task):
79
- """بازگرداندن نوع وظیفه براساس مدل."""
80
  task_map = {
81
  "translation": "translation",
82
  "qa": "question-answering",
83
  "persian_nlp": "text-classification",
84
  "custom_ai": "text-generation",
85
- "math": "text-generation"
86
  }
87
  return task_map.get(task, "text-generation")
88
 
89
- def process_task(self, task, **kwargs):
90
- """پردازش وظایف به صورت غیرهمزمان."""
91
  if task not in self.models:
92
- raise ValueError(f"Model for task '{task}' not loaded.")
93
-
94
- def task_handler():
95
- try:
96
- if task == "math":
97
- result = self.process_math_task(kwargs.get("text"))
98
- else:
99
- model = self.load_model_for_task(task)
100
- if model:
101
- result = model(**kwargs)
102
- else:
103
- result = None
104
- self.queue.put(result)
105
- except Exception as e:
106
- print(f"Error processing task '{task}': {e}")
107
- self.queue.put(None)
108
- finally:
109
- gc.collect() # آزادسازی حافظه پس از هر وظیفه
110
 
111
- self.executor.submit(task_handler)
 
 
 
 
112
 
113
  def process_math_task(self, text):
114
- """پردازش ریاضی با OpenAI."""
115
  try:
116
- MemoryManager.check_memory_usage()
117
- response = openai.Completion.create(
118
- engine="text-davinci-003",
119
- prompt=text,
120
- max_tokens=100
121
  )
122
- return response.choices[0].text.strip()
123
  except Exception as e:
124
- print(f"Error processing math task with OpenAI: {e}")
125
  return None
126
 
127
- def get_result(self):
128
- """دریافت نتیجه از صف."""
129
- return self.queue.get()
 
130
 
131
  # نمونه استفاده
132
  if __name__ == "__main__":
 
 
 
 
 
 
 
 
133
  system = MultiModelSystem()
134
 
135
- # وظایف مختلف
136
  tasks = [
137
- {"task": "translation", "kwargs": {"text": "سلام دنیا!", "src_lang": "fa", "tgt_lang": "en"}},
138
- {"task": "qa", "kwargs": {"question": "پایتخت ایران چیست؟", "context": "ایران کشوری در خاورمیانه است و پایتخت آن تهران است."}},
139
- {"task": "math", "kwargs": {"text": "What is the integral of x^2?"}},
140
- {"task": "persian_nlp", "kwargs": {"text": "این یک جمله فارسی است."}},
141
- {"task": "custom_ai", "kwargs": {"text": "تحلیل متن‌های تاریخی طبری."}}
142
  ]
143
 
144
- # ارسال وظایف برای پردازش
145
  for task_info in tasks:
146
- system.process_task(task_info["task"], **task_info["kwargs"])
147
-
148
- # بازیابی نتایج
149
- for _ in range(len(tasks)):
150
- print("Result:", system.get_result())
 
 
1
  from transformers import pipeline
2
+ import os
 
3
  import openai
4
+ import gc # برای آزاد کردن حافظه
5
+ import psutil # برای مانیتور کردن حافظه
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
 
7
  class MultiModelSystem:
8
  def __init__(self):
 
9
  self.models = {}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
10
 
11
+ def load_model(self, task, model_id, use_auth_token=None):
12
+ """مدل‌ها را به صورت lazy بارگذاری می‌کند."""
13
+ if model_id == "OpenAI":
14
+ return self.load_openai_model()
15
+ else:
 
 
16
  return pipeline(
17
  task=self.get_task_type(task),
18
  model=model_id,
19
+ use_auth_token=use_auth_token
 
20
  )
21
+
22
+ def unload_model(self, task):
23
+ """مدل بارگذاری شده را از حافظه پاک می‌کند."""
24
+ if task in self.models:
25
+ del self.models[task]
26
+ gc.collect() # جمع‌آوری حافظه
27
+
28
+ @staticmethod
29
+ def load_openai_model():
30
+ """مدل ریاضی OpenAI."""
31
+ return "OpenAI (Math)"
32
 
33
  @staticmethod
34
  def get_task_type(task):
 
35
  task_map = {
36
  "translation": "translation",
37
  "qa": "question-answering",
38
  "persian_nlp": "text-classification",
39
  "custom_ai": "text-generation",
40
+ "math": "text-generation",
41
  }
42
  return task_map.get(task, "text-generation")
43
 
44
+ def process_task(self, task, model_id, **kwargs):
45
+ """مدیریت وظایف."""
46
  if task not in self.models:
47
+ self.models[task] = self.load_model(task, model_id, use_auth_token=os.getenv("Passsssssss"))
48
+ model = self.models[task]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
49
 
50
+ # پردازش وظیفه
51
+ if task == "math":
52
+ return self.process_math_task(kwargs.get("text"))
53
+ else:
54
+ return model(**kwargs)
55
 
56
  def process_math_task(self, text):
57
+ """مدیریت وظایف ریاضی OpenAI."""
58
  try:
59
+ response = openai.ChatCompletion.create(
60
+ model="gpt-4",
61
+ messages=[{"role": "user", "content": text}]
 
 
62
  )
63
+ return response['choices'][0]['message']['content'].strip()
64
  except Exception as e:
65
+ print(f"Error processing math task: {e}")
66
  return None
67
 
68
+ # مانیتورینگ حافظه
69
+ def check_memory_usage():
70
+ mem = psutil.virtual_memory()
71
+ print(f"Memory usage: {mem.percent}% ({mem.used / (1024 ** 3):.2f} GB used)")
72
 
73
  # نمونه استفاده
74
  if __name__ == "__main__":
75
+ MODEL_CONFIG = {
76
+ "translation": "PontifexMaximus/opus-mt-iir-en-finetuned-fa-to-en",
77
+ "qa": "HooshvareLab/bert-fa-base-uncased",
78
+ "math": "OpenAI",
79
+ "persian_nlp": "HooshvareLab/bert-fa-zwnj-base",
80
+ "custom_ai": "universitytehran/PersianMind-v1.0",
81
+ }
82
+
83
  system = MultiModelSystem()
84
 
 
85
  tasks = [
86
+ {"task": "translation", "model_id": MODEL_CONFIG["translation"], "kwargs": {"text": "سلام دنیا!", "src_lang": "fa", "tgt_lang": "en"}},
87
+ {"task": "qa", "model_id": MODEL_CONFIG["qa"], "kwargs": {"question": "پایتخت ایران چیست؟", "context": "ایران کشوری در خاورمیانه است و پایتخت آن تهران است."}},
88
+ {"task": "math", "model_id": MODEL_CONFIG["math"], "kwargs": {"text": "What is the integral of x^2?"}},
 
 
89
  ]
90
 
 
91
  for task_info in tasks:
92
+ check_memory_usage() # نمایش میزان حافظه
93
+ result = system.process_task(task_info["task"], task_info["model_id"], **task_info["kwargs"])
94
+ print(f"Result for task '{task_info['task']}':", result)
95
+ system.unload_model(task_info["task"]) # آزادسازی حافظه
96
+ check_memory_usage() # دوباره بررسی حافظه