test / main.py
Guiziii's picture
Update main.py
d240a47 verified
raw
history blame contribute delete
805 Bytes
from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
app = Flask(__name__)
@app.route('/translate', methods=['POST'])
def translate():
# Load the model and tokenizer for each request
tokenizer = AutoTokenizer.from_pretrained("unicamp-dl/translation-en-pt-t5")
model = AutoModelForSeq2SeqLM.from_pretrained("unicamp-dl/translation-en-pt-t5")
enpt_pipeline = pipeline('text2text-generation', model=model, tokenizer=tokenizer, max_length=512, num_beams=5)
content = request.json
text_to_translate = content['text']
result = enpt_pipeline(f"translate English to Portuguese: {text_to_translate}")
return jsonify({"translated_text": result[0]["generated_text"]})
if __name__ == '__main__':
app.run(debug=True)