Spaces:
Runtime error
Runtime error
File size: 2,671 Bytes
c23173c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F
class TextEncoder:
def __init__(
self,
model_name: str = "nlptown/bert-base-multilingual-uncased-sentiment",
#model_name: str = "alec228/audio-sentiment/tree/main/bert-sentiment",
#cache_dir: str = "./models"
):
# Tokenizer pour prétraiter le texte
self.tokenizer = AutoTokenizer.from_pretrained(
model_name,
#cache_dir=cache_dir
)
# Modèle BERT de base (sans tête de classification)
self.model = AutoModel.from_pretrained(
model_name,
#cache_dir=cache_dir
)
def extract_features(self, texts: list[str]) -> torch.Tensor:
"""
Prend en entrée une liste de chaînes et renvoie
les embeddings du token [CLS] pour chaque texte.
"""
# 1. Tokenisation
inputs = self.tokenizer(
texts,
return_tensors="pt",
truncation=True,
padding=True
)
# 2. Passage dans le modèle sans calcul de gradient
with torch.no_grad():
outputs = self.model(**inputs)
# 3. Extraction de l'embedding du token [CLS]
return outputs.last_hidden_state[:, 0, :] # [batch, hidden_size]
def analyze_sentiment(text: str) -> dict:
"""
Analyse le sentiment d'un texte avec un modèle déjà fine-tuned
(nlptown/bert-base-multilingual-uncased-sentiment) et renvoie
un dict {label: confidence}.
"""
# Chargement du tokenizer et du modèle de classification
tokenizer = AutoTokenizer.from_pretrained(
"nlptown/bert-base-multilingual-uncased-sentiment",
#cache_dir="./models"
)
model = AutoModelForSequenceClassification.from_pretrained(
"nlptown/bert-base-multilingual-uncased-sentiment",
#cache_dir="./models"
)
# Préparation
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
probs = F.softmax(logits, dim=1).squeeze().tolist()
# Les classes vont de 1 à 5, on choisit la plus probable
label_idx = int(torch.argmax(torch.tensor(probs))) + 1
if label_idx <= 2:
label = "négatif"
elif label_idx == 3:
label = "neutre"
else:
label = "positif"
confidence = round(max(probs), 3)
return {label: confidence}
|