from flask import Flask, request, jsonify
from transformers import MT5ForConditionalGeneration, T5Tokenizer

app = Flask(__name__)

model = MT5ForConditionalGeneration.from_pretrained('iliemihai/mt5-base-romanian-diacritics', cache_dir='cache/')
tokenizer = T5Tokenizer.from_pretrained('iliemihai/mt5-base-romanian-diacritics', legacy=False, cache_dir='cache/')

@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})