jgrosjean-mathesis/sentence-swissbert

import torch
from transformers import AutoModel, AutoTokenizer

# Cargar swissBERT para incrustaciones de oraciones
model_name = "jgrosjean-mathesis/sentence-swissbert"
model = AutoModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def generate_sentence_embedding(sentence, language):
    # Configurar adaptador al idioma especificado
    if "de" in language:
        model.set_default_language("de_CH")
    if "fr" in language:
        model.set_default_language("fr_CH")
    if "it" in language:
        model.set_default_language("it_CH")
    if "rm" in language:
        model.set_default_language("rm_CH")

    # Tokenizar la oración de entrada
    inputs = tokenizer(sentence, padding=True, truncation=True, return_tensors="pt", max_length=512)

    # Pasar la entrada tokenizada por el modelo
    with torch.no_grad():
        outputs = model(**inputs)

    # Extraer incrustaciones de la oración mediante agrupación media
    token_embeddings = outputs.last_hidden_state
    attention_mask = inputs['attention_mask'].unsqueeze(-1).expand(token_embeddings.size()).float()
    sum_embeddings = torch.sum(token_embeddings * attention_mask, 1)
    sum_mask = torch.clamp(attention_mask.sum(1), min=1e-9)
    embedding = sum_embeddings / sum_mask

    return embedding

# Probarlo
sentence_0 = "Wir feiern am 1. August den Schweizer Nationalfeiertag."
sentence_0_embedding = generate_sentence_embedding(sentence_0, language="de")
print(sentence_0_embedding)

from sklearn.metrics.pairwise import cosine_similarity

# Definir dos oraciones
sentence_1 = ["Der Zug kommt um 9 Uhr in Zürich an."]
sentence_2 = ["Le train arrive à Lausanne à 9h."]

# Calcular incrustaciones para ambas
embedding_1 = generate_sentence_embedding(sentence_1, language="de")
embedding_2 = generate_sentence_embedding(sentence_2, language="fr")

# Calcular la similitud de coseno
cosine_score = cosine_similarity(embedding_1, embedding_2)

# Mostrar el puntaje
print("El puntaje de coseno para", sentence_1, "y", sentence_2, "es", cosine_score)