edyrkaj's picture
Update app.py support translation
408936f verified
raw
history blame
1.56 kB
from transformers.utils import logging
from language_directions import *
logging.set_verbosity_error()
import warnings
warnings.filterwarnings("ignore",
message="Using the model-agnostic default `max_length`")
import os
import gradio as gr
from transformers import pipeline
pipe = pipeline("image-to-text",
model="Salesforce/blip-image-captioning-base")
def translate(input_text, source, target):
source_readable = source
if source == "Auto Detect" or source.startswith("Detected"):
source, _ = auto_detect_language_code(input_text)
if source in source_lang_dict.keys():
source = source_lang_dict[source]
target_lang_dict, _ = get_target_languages(source)
try:
target = target_lang_dict[target]
# Helsinki-NLP/opus-mt-en-sq
model = f"Helsinki-NLP/opus-mt-{source}-{target}"
pipe = pipeline("translation", model=model)
translation = pipe(input_text)
return translation[0]['translation_text'], ""
except KeyError:
return "", f"Error: Translation direction {source_readable} to {target} is not supported by Helsinki Translation Models"
def launch(input):
out = pipe(input)
context_str = out[0]['generated_text']
translate_str = translate(context_str, 'en', 'sq')
return translate_str
iface = gr.Interface(launch,
inputs=gr.Image(type='pil'),
outputs="text")
iface.launch(share=True)
# iface.launch(share=True,
# server_port=int(os.environ['PORT1']))
iface.close()