import torch
import os
from transformers.pipelines import pipeline
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
from flask import Flask, request
from flask_restful import Api, Resource
from flask_cors import CORS

app = Flask(__name__)
CORS(app)
cors = CORS(app, resource={
    r"/*": {
        "origins": "*"
    }
})
api = Api(app)

app.config['CORS_HEADERS'] = 'Content-Type'


class Classifier():

    def __init__(self, data_en):
        self.model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-one-to-many-mmt")
        self.tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-one-to-many-mmt", src_lang="en_XX")
        self.model_inputs = self.tokenizer(data_en, return_tensors="pt")

        # translate from English to Malayalam
        self.generated_tokens = self.model.generate(
            **self.model_inputs,
            forced_bos_token_id=self.tokenizer.lang_code_to_id["ml_IN"]
        )
        self. translate = self.tokenizer.batch_decode(self.generated_tokens, skip_special_tokens=True)
        self.data_en = data_en

    def get_translator(self):
        output = self.translate(self.data_en)
        return {'output': output}


class Translate(Resource):
    def post(self):
        try:
            # Decode json object from the request
            json_object = request.get_json()
            data_en = json_object["text"]
            obj = Classifier(data_en)
        except Exception as e:
            return {"Message": "Error in creating Translator object" + str(e)}
        status = obj.get_translator()
        return status


api.add_resource(Translate, '/api/translate')

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)