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.")