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/alephbert-base")
english_model = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.3")

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

    # Detect language 
    language = detect(user_input)

    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(user_input, 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)