darijapfe / app.py
AN2002's picture
Update app.py
1acd54a verified
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()