from flask import Flask, request, jsonify
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

app = Flask(__name__)

# Load your model (can be any HF model)
MODEL_NAME = "tiiuae/falcon-7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

@app.route('/setToken', methods=['POST'])
def set_token():
    """
    Main multimodal API endpoint.
    Handles system + user prompts and returns generated response.
    """
    data = request.get_json(force=True)

    system_prompt = data.get("system_prompt", "You are a helpful AI.")
    user_input = data.get("user_input", "")
    temperature = float(data.get("temperature", 0.7))
    mode = data.get("mode", "text")

    # Text mode (default)
    if mode == "text":
        full_prompt = f"{system_prompt}\nUser: {user_input}\nAI:"
        inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=True,
            temperature=temperature
        )
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        response = response.split("AI:")[-1].strip()
        return jsonify({
            "model": MODEL_NAME,
            "response": response,
            "mode": "text"
        })

    # You can later add multimodal branches here:
    # - "image" -> call image generation pipeline
    # - "audio" -> call speech-to-text / text-to-speech
    # - "embedding" -> return vector embeddings

    return jsonify({"error": f"Unsupported mode: {mode}"}), 400

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