translation_api.py

from flask import jsonify, send_file, request
from main import *

def perform_translation(text, target_language_code='es_XX', source_language_code='en_XX'):
    if translation_model is None:
        return {"error": "Translation model not initialized."}

    encoded_text = translation_model.tokenizer(text, return_tensors="pt", padding=True).to(device)
    generated_tokens = translation_model.generate(input_ids=encoded_text['input_ids'], attention_mask=encoded_text['attention_mask'], forced_bos_token_id=translation_model.config.lang_code_to_id[target_language_code])
    translation = translation_model.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]

    return {"translated_text": translation}

def translation_api():
    data = request.get_json()
    text = data.get('text')
    target_lang = data.get('target_lang', 'es')
    source_lang = data.get('source_lang', 'en')
    if not text:
        return jsonify({"error": "Text is required"}), 400
    output = perform_translation(text, target_language_code=f'{target_lang}_XX', source_language_code=f'{source_lang}_XX')
    if "error" in output:
        return jsonify({"error": output["error"]}), 500
    return jsonify(output)