darkbert / darkbert_model /analyse_darkbert.py
Bartholomew-1's picture
Upload 4 files
5329bd8 verified
Raw
History Blame Contribute Delete
1.34 kB
from transformers import AutoTokenizer, AutoModel
import torch
import os
from sklearn.metrics.pairwise import cosine_similarity
# Initialisation de DarkBERT
tokenizer = AutoTokenizer.from_pretrained("kakaobrain/DarkBERT")
model = AutoModel.from_pretrained("kakaobrain/DarkBERT").to("cuda")
def analyser_fichier(chemin):
with open(chemin, "r", encoding="utf-8") as f:
texte = f.read()
inputs = tokenizer(texte, return_tensors="pt", truncation=True, max_length=512).to("cuda")
with torch.no_grad():
sorties = model(**inputs)
return sorties.last_hidden_state.mean(dim=1).cpu().numpy()
# Parcours des fichiers sur le bureau
bureau = "C:/Users/Utilisateur/Desktop"
fichiers = [f for f in os.listdir(bureau) if f.endswith(('.txt', '.csv'))][:2]
resultats = {}
for fichier in fichiers:
chemin = os.path.join(bureau, fichier)
resultats[fichier] = analyser_fichier(chemin)
print(f"Fichier {fichier} traité - Dimension embedding: {resultats[fichier].shape}")
# Calcul de la similarité sémantique
if len(fichiers) == 2:
similarite = cosine_similarity(resultats[fichiers[0]], resultats[fichiers[1]])[0][0]
print(f"\nSimilarité sémantique entre les fichiers : {similarite:.2f}")
else:
print("\nMoins de deux fichiers trouvés pour comparer.")