Batnini commited on
Commit
65b1fd2
·
verified ·
1 Parent(s): 47c9040

Update tools/arabic_generator.py

Browse files
Files changed (1) hide show
  1. tools/arabic_generator.py +25 -1
tools/arabic_generator.py CHANGED
@@ -1 +1,25 @@
1
- This is a dummy file.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+
3
+ class ArabicTextGenerator:
4
+ def __init__(self):
5
+ self.model_loaded = False
6
+
7
+ def load_model(self):
8
+ """Lazy load the model to improve startup time"""
9
+ if not self.model_loaded:
10
+ self.generator = pipeline(
11
+ 'text-generation',
12
+ model='aubmindlab/aragpt2-base',
13
+ device=-1 # Force CPU
14
+ )
15
+ self.model_loaded = True
16
+
17
+ def generate(self, prompt, max_length=100):
18
+ self.load_model()
19
+ output = self.generator(
20
+ prompt,
21
+ max_length=max_length,
22
+ num_return_sequences=1,
23
+ pad_token_id=50256 # eos_token_id for aragpt2
24
+ )
25
+ return output[0]['generated_text']