# -*- coding: utf-8 -*-
import torch
from flask import Flask, request, jsonify, render_template
from transformers import AutoTokenizer, AutoModel

app = Flask(__name__)

# Chargez le modèle et le tokenizer
tokenizer = AutoTokenizer.from_pretrained("/mnt/e/1__IA/2_Modeles/1__Falcon/7B")
model = AutoModel.from_pretrained("/mnt/e/1__IA/2_Modeles/1__Falcon/7B")

# Convertir le modèle à l'API BetterTransformer
model = model.to_bettertransformer()

# Déplacez le modèle vers le CPU
model.to("cpu")

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    # Obtenez le texte de la requête
    data = request.get_json()
    text = data['text']

    # Tokenisez et prédisez
    inputs = tokenizer(text, return_tensors="pt")
    
    # Supprimez l'argument position_ids
    if 'position_ids' in inputs:
        del inputs['position_ids']

    outputs = model(**inputs)

    # Libérer la mémoire GPU inutilisée
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Ici, nous renvoyons simplement la forme des sorties pour la démonstration
    return jsonify({'shape': str(outputs.last_hidden_state.shape)})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=False)