from flask import Flask, request, jsonify
from transformers import pipeline
from langdetect import detect

app = Flask(__name__)

# load models
hebrew_model = pipeline("text-generation", model="onlplab/alephbertgpt") 
english_model = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct")

@app.route("/generate", methods=["POST"])
def generate():
    data = request.json
    text = data.get("text", "")

    # detect language 
    language = detect(text)

    if language == 'he':
        model = hebrew_model
    elif language == 'en':
        model = english_model
    else:
        print("Decision Making Helper BOT currently supports Hebrew and English Languages")

    # create an answer from the model
    response = model(text, max_length=100, do_sample=True)
    return jsonify({"response": response[0]['generated_text']})

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