EGYADMIN commited on
Commit
171f954
·
verified ·
1 Parent(s): a3b5c66

Create models/model_loader.py

Browse files
Files changed (1) hide show
  1. models/model_loader.py +327 -1
models/model_loader.py CHANGED
@@ -24,4 +24,330 @@ class ModelLoader:
24
  المعاملات:
25
  ----------
26
  config : Dict, optional
27
- إ
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
  المعاملات:
25
  ----------
26
  config : Dict, optional
27
+ إعدادات محمّل النماذج
28
+ use_gpu : bool, optional
29
+ استخدام GPU إذا كان متاحًا
30
+ """
31
+ self.config = config or {}
32
+ self.use_gpu = use_gpu and torch.cuda.is_available()
33
+
34
+ # التحقق من توفر GPU
35
+ if self.use_gpu:
36
+ self.device = torch.device("cuda")
37
+ logger.info(f"تم اكتشاف GPU: {torch.cuda.get_device_name(0)}")
38
+ logger.info(f"ذاكرة GPU المتاحة: {self._get_available_gpu_memory()} ميجابايت")
39
+ else:
40
+ self.device = torch.device("cpu")
41
+ logger.info("استخدام المعالج المركزي CPU")
42
+
43
+ # تهيئة قواميس لتخزين النماذج المحملة
44
+ self.models = {}
45
+ self.tokenizers = {}
46
+
47
+ # تحميل النماذج المطلوبة بشكل افتراضي
48
+ self._load_default_models()
49
+
50
+ logger.info("تم تهيئة محمّل النماذج")
51
+
52
+ def get_ner_model(self):
53
+ """
54
+ الحصول على نموذج التعرف على الكيانات المسماة
55
+
56
+ المخرجات:
57
+ --------
58
+ أي
59
+ نموذج التعرف على الكيانات المسماة
60
+ """
61
+ if "ner" in self.models:
62
+ return self.models["ner"]
63
+
64
+ # تحميل النموذج إذا لم يكن محملاً
65
+ ner_model = self._load_ner_model()
66
+
67
+ return ner_model
68
+
69
+ def get_similarity_model(self):
70
+ """
71
+ الحصول على نموذج التشابه النصي
72
+
73
+ المخرجات:
74
+ --------
75
+ أي
76
+ نموذج التشابه النصي
77
+ """
78
+ if "similarity" in self.models:
79
+ return self.models["similarity"]
80
+
81
+ # تحميل النموذج إذا لم يكن محملاً
82
+ similarity_model = self._load_similarity_model()
83
+
84
+ return similarity_model
85
+
86
+ def get_classification_model(self):
87
+ """
88
+ الحصول على نموذج التصنيف
89
+
90
+ المخرجات:
91
+ --------
92
+ أي
93
+ نموذج التصنيف
94
+ """
95
+ if "classification" in self.models:
96
+ return self.models["classification"]
97
+
98
+ # تحميل النموذج إذا لم يكن محملاً
99
+ classification_model = self._load_classification_model()
100
+
101
+ return classification_model
102
+
103
+ def get_tokenizer(self, model_name):
104
+ """
105
+ الحصول على محلل الترميز للنموذج
106
+
107
+ المعاملات:
108
+ ----------
109
+ model_name : str
110
+ اسم النموذج
111
+
112
+ المخرجات:
113
+ --------
114
+ أي
115
+ محلل الترميز
116
+ """
117
+ if model_name in self.tokenizers:
118
+ return self.tokenizers[model_name]
119
+
120
+ # تحميل محلل الترميز إذا لم يكن محملاً
121
+ if model_name == "ner":
122
+ tokenizer = self._load_tokenizer("aubmindlab/bert-base-arabertv02-ner")
123
+ elif model_name == "similarity":
124
+ tokenizer = self._load_tokenizer("UBC-NLP/ARBERT")
125
+ elif model_name == "classification":
126
+ tokenizer = self._load_tokenizer("CAMeL-Lab/bert-base-arabic-camelbert-mix")
127
+ else:
128
+ raise ValueError(f"محلل الترميز غير معروف: {model_name}")
129
+
130
+ self.tokenizers[model_name] = tokenizer
131
+ return tokenizer
132
+
133
+ def release_model(self, model_name):
134
+ """
135
+ تحرير النموذج من الذاكرة
136
+
137
+ المعاملات:
138
+ ----------
139
+ model_name : str
140
+ اسم النموذج
141
+ """
142
+ if model_name in self.models:
143
+ del self.models[model_name]
144
+
145
+ if self.use_gpu:
146
+ torch.cuda.empty_cache()
147
+ gc.collect()
148
+
149
+ logger.info(f"تم تحرير النموذج: {model_name}")
150
+
151
+ def release_all_models(self):
152
+ """
153
+ تحرير جميع النماذج من الذاكرة
154
+ """
155
+ self.models = {}
156
+ self.tokenizers = {}
157
+
158
+ if self.use_gpu:
159
+ torch.cuda.empty_cache()
160
+ gc.collect()
161
+
162
+ logger.info("تم تحرير جميع النماذج")
163
+
164
+ def get_available_memory(self):
165
+ """
166
+ الحصول على كمية الذاكرة المتاحة
167
+
168
+ المخرجات:
169
+ --------
170
+ int
171
+ كمية الذاكرة المتاحة بالميجابايت
172
+ """
173
+ if self.use_gpu:
174
+ return self._get_available_gpu_memory()
175
+ else:
176
+ # ليس هناك طريقة موحدة للحصول على ذاكرة CPU
177
+ return 0
178
+
179
+ def _load_default_models(self):
180
+ """
181
+ تحميل النماذج الافتراضية
182
+ """
183
+ default_models = self.config.get("default_models", [])
184
+
185
+ for model_name in default_models:
186
+ try:
187
+ if model_name == "ner":
188
+ self._load_ner_model()
189
+ elif model_name == "similarity":
190
+ self._load_similarity_model()
191
+ elif model_name == "classification":
192
+ self._load_classification_model()
193
+ else:
194
+ logger.warning(f"نموذج غير معروف: {model_name}")
195
+ except Exception as e:
196
+ logger.error(f"فشل في تحميل النموذج {model_name}: {str(e)}")
197
+
198
+ def _load_ner_model(self):
199
+ """
200
+ تحميل نموذج التعرف على الكيانات المسماة
201
+
202
+ المخرجات:
203
+ --------
204
+ أي
205
+ نموذج التعرف على الكيانات المسماة
206
+ """
207
+ try:
208
+ from transformers import pipeline
209
+
210
+ logger.info("جاري تحميل نموذج التعرف على الكيانات المسماة...")
211
+
212
+ model_name = self.config.get("ner_model", "aubmindlab/bert-base-arabertv02-ner")
213
+
214
+ # تحميل النموذج
215
+ ner_model = pipeline(
216
+ "token-classification",
217
+ model=model_name,
218
+ aggregation_strategy="simple",
219
+ device=0 if self.use_gpu else -1
220
+ )
221
+
222
+ self.models["ner"] = ner_model
223
+ logger.info(f"تم تحميل نموذج التعرف على الكيانات المسماة: {model_name}")
224
+
225
+ return ner_model
226
+
227
+ except Exception as e:
228
+ logger.error(f"فشل في تحميل نموذج التعرف على الكيانات المسماة: {str(e)}")
229
+ raise
230
+
231
+ def _load_similarity_model(self):
232
+ """
233
+ تحميل نموذج التشابه النصي
234
+
235
+ المخرجات:
236
+ --------
237
+ أي
238
+ نموذج التشابه النصي
239
+ """
240
+ try:
241
+ from transformers import AutoModel
242
+
243
+ logger.info("جاري تحميل نموذج التشابه النصي...")
244
+
245
+ model_name = self.config.get("similarity_model", "UBC-NLP/ARBERT")
246
+
247
+ # تحميل النموذج
248
+ similarity_model = AutoModel.from_pretrained(model_name)
249
+
250
+ if self.use_gpu:
251
+ similarity_model = similarity_model.to(self.device)
252
+
253
+ self.models["similarity"] = similarity_model
254
+ logger.info(f"تم تحميل نموذج التشابه النصي: {model_name}")
255
+
256
+ return similarity_model
257
+
258
+ except Exception as e:
259
+ logger.error(f"فشل في تحميل نموذج التشابه النصي: {str(e)}")
260
+ raise
261
+
262
+ def _load_classification_model(self):
263
+ """
264
+ تحميل نموذج التصنيف
265
+
266
+ المخرجات:
267
+ --------
268
+ أي
269
+ نموذج التصنيف
270
+ """
271
+ try:
272
+ from transformers import pipeline
273
+
274
+ logger.info("جاري تحميل نموذج التصنيف...")
275
+
276
+ model_name = self.config.get("classification_model", "CAMeL-Lab/bert-base-arabic-camelbert-mix")
277
+
278
+ # تحميل النموذج
279
+ classification_model = pipeline(
280
+ "text-classification",
281
+ model=model_name,
282
+ device=0 if self.use_gpu else -1
283
+ )
284
+
285
+ self.models["classification"] = classification_model
286
+ logger.info(f"تم تحميل نموذج التصنيف: {model_name}")
287
+
288
+ return classification_model
289
+
290
+ except Exception as e:
291
+ logger.error(f"فشل في تحميل نموذج التصنيف: {str(e)}")
292
+ raise
293
+
294
+ def _load_tokenizer(self, model_name):
295
+ """
296
+ تحميل محلل الترميز
297
+
298
+ المعاملات:
299
+ ----------
300
+ model_name : str
301
+ اسم النموذج
302
+
303
+ المخرجات:
304
+ --------
305
+ أي
306
+ محلل الترميز
307
+ """
308
+ try:
309
+ from transformers import AutoTokenizer
310
+
311
+ logger.info(f"جاري تحميل محلل الترميز: {model_name}...")
312
+
313
+ # تحميل محلل الترميز
314
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
315
+
316
+ logger.info(f"تم تحميل محلل الترميز: {model_name}")
317
+
318
+ return tokenizer
319
+
320
+ except Exception as e:
321
+ logger.error(f"فشل في تحميل محلل الترميز {model_name}: {str(e)}")
322
+ raise
323
+
324
+ def _get_available_gpu_memory(self):
325
+ """
326
+ الحصول على كمية ذاكرة GPU المتاحة
327
+
328
+ المخرجات:
329
+ --------
330
+ int
331
+ كمية ذاكرة GPU المتاحة بالميجابايت
332
+ """
333
+ if not self.use_gpu:
334
+ return 0
335
+
336
+ try:
337
+ torch.cuda.empty_cache()
338
+ # الحصول على إجمالي ذاكرة GPU
339
+ total_memory = torch.cuda.get_device_properties(0).total_memory
340
+ # الحصول على ذاكرة GPU المستخدمة
341
+ allocated_memory = torch.cuda.memory_allocated(0)
342
+ # الحصول على ذاكرة GPU المحجوزة
343
+ reserved_memory = torch.cuda.memory_reserved(0)
344
+
345
+ # حساب الذاكرة المتاحة
346
+ available_memory = total_memory - allocated_memory - reserved_memory
347
+
348
+ # تحويل إلى ميجابايت
349
+ return available_memory / (1024 * 1024)
350
+
351
+ except Exception as e:
352
+ logger.error(f"فشل في الحصول على ذاكرة GPU المتاحة: {str(e)}")
353
+ return 0