dia-critic / api /index.py
Florin Bobiș
changes
c5a40da
raw
history blame
917 Bytes
from flask import Flask, request, jsonify
from transformers import MT5ForConditionalGeneration, T5Tokenizer
app = Flask(__name__)
model = MT5ForConditionalGeneration.from_pretrained('iliemihai/mt5-base-romanian-diacritics')
tokenizer = T5Tokenizer.from_pretrained('iliemihai/mt5-base-romanian-diacritics', legacy=False)
@app.route('/generate', methods=['POST'])
def generate_text():
input_text = request.get_json()['input_text']
if input_text is None:
return jsonify({'error': 'No input text provided'})
inputs = tokenizer(input_text, max_length=256, truncation=True, return_tensors="pt")
outputs = model.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"])
output = tokenizer.decode(outputs[0], skip_special_tokens=True)
return jsonify({'output': output})
# if __name__ == '__main__':
# app.run(host="0.0.0.0", debug=True, port=8082, threaded=True)