edyrkaj commited on
Commit
408936f
1 Parent(s): c16d8e2

Update app.py support translation

Browse files
Files changed (1) hide show
  1. app.py +22 -1
app.py CHANGED
@@ -1,4 +1,5 @@
1
  from transformers.utils import logging
 
2
  logging.set_verbosity_error()
3
 
4
  import warnings
@@ -12,9 +13,29 @@ from transformers import pipeline
12
  pipe = pipeline("image-to-text",
13
  model="Salesforce/blip-image-captioning-base")
14
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
  def launch(input):
16
  out = pipe(input)
17
- return out[0]['generated_text']
 
 
18
 
19
  iface = gr.Interface(launch,
20
  inputs=gr.Image(type='pil'),
 
1
  from transformers.utils import logging
2
+ from language_directions import *
3
  logging.set_verbosity_error()
4
 
5
  import warnings
 
13
  pipe = pipeline("image-to-text",
14
  model="Salesforce/blip-image-captioning-base")
15
 
16
+
17
+ def translate(input_text, source, target):
18
+ source_readable = source
19
+ if source == "Auto Detect" or source.startswith("Detected"):
20
+ source, _ = auto_detect_language_code(input_text)
21
+ if source in source_lang_dict.keys():
22
+ source = source_lang_dict[source]
23
+ target_lang_dict, _ = get_target_languages(source)
24
+ try:
25
+ target = target_lang_dict[target]
26
+ # Helsinki-NLP/opus-mt-en-sq
27
+ model = f"Helsinki-NLP/opus-mt-{source}-{target}"
28
+ pipe = pipeline("translation", model=model)
29
+ translation = pipe(input_text)
30
+ return translation[0]['translation_text'], ""
31
+ except KeyError:
32
+ return "", f"Error: Translation direction {source_readable} to {target} is not supported by Helsinki Translation Models"
33
+
34
  def launch(input):
35
  out = pipe(input)
36
+ context_str = out[0]['generated_text']
37
+ translate_str = translate(context_str, 'en', 'sq')
38
+ return translate_str
39
 
40
  iface = gr.Interface(launch,
41
  inputs=gr.Image(type='pil'),