fix: slow model initialization
Browse files- frasaria/__init__.py +15 -3
frasaria/__init__.py
CHANGED
@@ -1,6 +1,18 @@
|
|
1 |
-
|
|
|
2 |
from parrot import Parrot
|
|
|
|
|
|
|
3 |
|
|
|
|
|
|
|
|
|
4 |
parrot = Parrot(model_tag='prithivida/parrot_paraphraser_on_T5')
|
5 |
-
|
6 |
-
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
from dotenv import load_dotenv
|
3 |
from parrot import Parrot
|
4 |
+
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
5 |
+
from sentence_transformers import SentenceTransformer
|
6 |
+
import deepl
|
7 |
|
8 |
+
load_dotenv()
|
9 |
+
DEEPL_API_KEY = os.getenv('DEEPL_API_KEY')
|
10 |
+
|
11 |
+
# initialize the models
|
12 |
parrot = Parrot(model_tag='prithivida/parrot_paraphraser_on_T5')
|
13 |
+
for model_tag in ['prithivida/parrot_fluency_model', 'prithivida/parrot_adequacy_model']:
|
14 |
+
AutoModelForSequenceClassification.from_pretrained(model_tag)
|
15 |
+
AutoTokenizer.from_pretrained(model_tag)
|
16 |
+
SentenceTransformer("sentence-transformers/paraphrase-distilroberta-base-v2")
|
17 |
+
|
18 |
+
translator = deepl.Translator(DEEPL_API_KEY)
|