Spaces:
Runtime error
Runtime error
File size: 4,693 Bytes
e69a4b4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
# import os
# import torch
# import gradio s gr
# import time
# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
# from flores200_codes import flores_codes
# def load_models():
# # build model and tokenizer
# model_name_dict = {'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M',
# #'nllb-1.3B': 'facebook/nllb-200-1.3B',
# #'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B',
# #'nllb-3.3B': 'facebook/nllb-200-3.3B',
# }
# model_dict = {}
# for call_name, real_name in model_name_dict.items():
# print('\tLoading model: %s' % call_name)
# model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
# tokenizer = AutoTokenizer.from_pretrained(real_name)
# model_dict[call_name+'_model'] = model
# model_dict[call_name+'_tokenizer'] = tokenizer
# return model_dict
# def translation(source, target, text):
# if len(model_dict) == 2:
# model_name = 'nllb-distilled-600M'
# start_time = time.time()
# source = flores_codes[source]
# target = flores_codes[target]
# model = model_dict[model_name + '_model']
# tokenizer = model_dict[model_name + '_tokenizer']
# translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source, tgt_lang=target)
# output = translator(text, max_length=400)
# end_time = time.time()
# output = output[0]['translation_text']
# result = {'inference_time': end_time - start_time,
# 'source': source,
# 'target': target,
# 'result': output}
# return result
# if __name__ == '__main__':
# print('\tinit models')
# global model_dict
# model_dict = load_models()
# # define gradio demo
# lang_codes = list(flores_codes.keys())
# #inputs = [gr.inputs.Radio(['nllb-distilled-600M', 'nllb-1.3B', 'nllb-distilled-1.3B'], label='NLLB Model'),
# inputs = [gr.inputs.Dropdown(lang_codes, default='English', label='Source'),
# gr.inputs.Dropdown(lang_codes, default='Korean', label='Target'),
# gr.inputs.Textbox(lines=5, label="Input text"),
# ]
# outputs = gr.outputs.JSON()
# title = "NLLB distilled 600M demo"
# demo_status = "Demo is running on CPU"
# description = f"Details: https://github.com/facebookresearch/fairseq/tree/nllb. {demo_status}"
# examples = [
# ['English', 'Korean', 'Hi. nice to meet you']
# ]
# gr.Interface(translation,
# inputs,
# outputs,
# title=title,
# description=description,
# ).launch()
import os
import time
from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
from flores200_codes import flores_codes
app = Flask(__name__)
def load_models():
model_name_dict = {'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M'}
model_dict = {}
for call_name, real_name in model_name_dict.items():
print(f'\tLoading model: {call_name}')
model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
tokenizer = AutoTokenizer.from_pretrained(real_name)
model_dict[call_name + '_model'] = model
model_dict[call_name + '_tokenizer'] = tokenizer
return model_dict
global model_dict
model_dict = load_models()
@app.route('/api/translate', methods=['POST'])
def translate_text():
data = request.json
source_lang = data.get('source')
target_lang = data.get('target')
input_text = data.get('text')
if not source_lang or not target_lang or not input_text:
return jsonify({"error": "source, target, and text fields are required"}), 400
model_name = 'nllb-distilled-600M'
start_time = time.time()
source = flores_codes.get(source_lang)
target = flores_codes.get(target_lang)
if not source or not target:
return jsonify({"error": "Invalid source or target language code"}), 400
model = model_dict[model_name + '_model']
tokenizer = model_dict[model_name + '_tokenizer']
translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source, tgt_lang=target)
output = translator(input_text, max_length=400)
end_time = time.time()
output_text = output[0]['translation_text']
result = {
'inference_time': end_time - start_time,
'source': source_lang,
'target': target_lang,
'result': output_text
}
return jsonify(result)
if __name__ == '__main__':
app.run(debug=True)
|