import gradio as gr
import joblib
import numpy as np
from transformers import BertTokenizer, BertModel
import torch

# Charger le modèle SVM et la normalisation
svm_model = joblib.load("svm_model5.joblib")
tokenizer = BertTokenizer.from_pretrained("SI2M-Lab/DarijaBERT")
bert_model = BertModel.from_pretrained("SI2M-Lab/DarijaBERT")

# Fonction pour transformer le texte en embeddings BERT
def get_bert_embedding(text):
    with torch.no_grad():
        inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=128)
        outputs = bert_model(**inputs)
        cls_embedding = outputs.last_hidden_state[:, 0, :].squeeze().cpu().numpy()
    return np.array(cls_embedding).reshape(1, -1)

# Fonction de prédiction
def predict(text):
    embedding = get_bert_embedding(text)
    prediction = svm_model.predict(embedding)[0]  # Prédiction de classe
    return f"Classe prédite : {prediction}"

# Interface Gradio
iface = gr.Interface(fn=predict, inputs="text", outputs="text", title="Classification de texte avec SVM & BERT")

# Lancer l'application
iface.launch()