Update app.py
Browse files
app.py
CHANGED
@@ -1,7 +1,7 @@
|
|
1 |
import gradio as gr
|
2 |
import spaces
|
3 |
import torch
|
4 |
-
from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
|
5 |
import languagecodes
|
6 |
|
7 |
favourite_langs = {"German": "de", "Romanian": "ro", "English": "en", "-----": "-----"}
|
@@ -14,8 +14,10 @@ models = ["Helsinki-NLP", "t5-base", "t5-small", "t5-large",
|
|
14 |
"facebook/nllb-200-distilled-600M",
|
15 |
"facebook/nllb-200-distilled-1.3B",
|
16 |
"facebook/mbart-large-50-many-to-many-mmt",
|
|
|
17 |
"Unbabel/TowerInstruct-7B-v0.2",
|
18 |
-
"Unbabel/TowerInstruct-Mistral-7B-v0.2"
|
|
|
19 |
|
20 |
def model_to_cuda(model):
|
21 |
# Move the model to GPU if available
|
@@ -26,6 +28,17 @@ def model_to_cuda(model):
|
|
26 |
print("CUDA not available! Using CPU.")
|
27 |
return model
|
28 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
29 |
@spaces.GPU
|
30 |
def translate_text(input_text, sselected_language, tselected_language, model_name):
|
31 |
sl = all_langs[sselected_language]
|
@@ -43,7 +56,10 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
|
|
43 |
model = model_to_cuda(AutoModelForSeq2SeqLM.from_pretrained(model_name))
|
44 |
except EnvironmentError as error:
|
45 |
return f"Error finding model: {model_name}! Try other available language combination.", error
|
46 |
-
|
|
|
|
|
|
|
47 |
if 'nllb' in model_name:
|
48 |
tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=languagecodes.nllb_language_codes[sselected_language])
|
49 |
model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")
|
|
|
1 |
import gradio as gr
|
2 |
import spaces
|
3 |
import torch
|
4 |
+
from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, pipeline
|
5 |
import languagecodes
|
6 |
|
7 |
favourite_langs = {"German": "de", "Romanian": "ro", "English": "en", "-----": "-----"}
|
|
|
14 |
"facebook/nllb-200-distilled-600M",
|
15 |
"facebook/nllb-200-distilled-1.3B",
|
16 |
"facebook/mbart-large-50-many-to-many-mmt",
|
17 |
+
"utter-project/EuroLLM-1.7B",
|
18 |
"Unbabel/TowerInstruct-7B-v0.2",
|
19 |
+
"Unbabel/TowerInstruct-Mistral-7B-v0.2"
|
20 |
+
]
|
21 |
|
22 |
def model_to_cuda(model):
|
23 |
# Move the model to GPU if available
|
|
|
28 |
print("CUDA not available! Using CPU.")
|
29 |
return model
|
30 |
|
31 |
+
def eurollm(model_id, sl, tl, input_text):
|
32 |
+
model_id = "utter-project/EuroLLM-1.7B"
|
33 |
+
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
34 |
+
model = AutoModelForCausalLM.from_pretrained(model_id)
|
35 |
+
prompt = f"{sl}: {input_text}. {tl}:"
|
36 |
+
inputs = tokenizer(prompt, return_tensors="pt")
|
37 |
+
outputs = model.generate(**inputs, max_new_tokens=20)
|
38 |
+
output = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
39 |
+
print(output)
|
40 |
+
return output
|
41 |
+
|
42 |
@spaces.GPU
|
43 |
def translate_text(input_text, sselected_language, tselected_language, model_name):
|
44 |
sl = all_langs[sselected_language]
|
|
|
56 |
model = model_to_cuda(AutoModelForSeq2SeqLM.from_pretrained(model_name))
|
57 |
except EnvironmentError as error:
|
58 |
return f"Error finding model: {model_name}! Try other available language combination.", error
|
59 |
+
if 'eurollm' in model_name:
|
60 |
+
translated_text = eurollm(model_name, sl, tl, input_text)
|
61 |
+
return translated_text[0]['translation_text'], message_text
|
62 |
+
|
63 |
if 'nllb' in model_name:
|
64 |
tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=languagecodes.nllb_language_codes[sselected_language])
|
65 |
model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")
|