from flask import Flask, jsonify, request, render_template
from transformers import AutoAdapterModel, AutoTokenizer, TextClassificationPipeline

tokenizer = AutoTokenizer.from_pretrained("UBC-NLP/MARBERT")
model = AutoAdapterModel.from_pretrained("UBC-NLP/MARBERT")

model.load_adapter("nehalelkaref/aoc3_adapter", set_active=True, with_head=False, source="hf")
model.load_adapter("nehalelkaref/aoc4_adapter", set_active=True, with_head=False, source="hf")
model.load_adapter("nehalelkaref/sarcasm_adapter", set_active=True, with_head=False, source="hf")

model.load_adapter_fusion("nehalelkaref/region_fusion",with_head=True, set_active=True, source="hf")

pipe = TextClassificationPipeline(tokenizer=tokenizer, model=model)

app = Flask(__name__)

@app.route("/", methods=['GET'])
def home():
    return render_template('home.html')

@app.route('/classify', methods = ['POST'])
def classify():
    text = request.json['inputs']

    prediction = pipe(text)
    labels = {"LABEL_0":"GULF", "LABEL_1":"LEVANT","LABEL_2":"EGYPT"}
    regions = []
    for res in prediction:
        regions.append(labels[res['label']])
    
    return render_template('prediction.html', output=regions[0])

if __name__ == "__main__":
  app.run()