measure text similarity with Roberta

transformer.py

!pip install sentence_transformers

from sentence_transformers import SentenceTransformer, util
# use roberta 
model = SentenceTransformer('stsb-roberta-large')

def create_heatmap(similarity, cmap = "YlGnBu"):
    df = pd.DataFrame(similarity)
    df.columns = ['john', 'luke','mark', 'matt'] #ohn 0 mark 2 matt 3 luke 1
    df.index = ['john', 'luke','mark', 'matt']
    fig, ax = plt.subplots(figsize=(5,5))
    sns.heatmap(df, cmap=cmap)

# encode the input text
embeddings = model.encode(data, convert_to_tensor=True)

similarity = []
for i in range(len(data)):
    row = []
    for j in range(len(data)):
        row.append(util.pytorch_cos_sim(embeddings[i], embeddings[j]).item())
    similarity.append(row)
    
create_heatmap(similarity)