Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Select an option

  • Save AnteaterKit/7587dd27d27d7202e5a310030109f894 to your computer and use it in GitHub Desktop.

Select an option

Save AnteaterKit/7587dd27d27d7202e5a310030109f894 to your computer and use it in GitHub Desktop.
Further training of the embedding model (nomic-ai/nomic-embed-text-v1.5) to improve separation of vector representations of cars of different brands using a custom loss function.
Визуализация эмбеддингов
Методы: PCA и t-SNE для снижения размерности до 2D/3D.
Библиотеки: matplotlib (статические графики), plotly (интерактивные 3D-графики).
Кастомная функция потерь BrandAwareTripletLoss
Цель: Усилить разделение между брендами.
Компоненты:
intra_loss: уменьшает расстояние внутри одного бренда.
inter_loss: увеличивает расстояние между разными брендами (классический triplet loss).
Коэффициенты:
intra_coef = 0.01 (слабое влияние).
inter_coef = 0.9 (сильный акцент на разделение брендов).
Данные для обучения
Триплеты: 58 примеров вида (anchor, positive, negative).
Структура:
anchor и positive – автомобили одного бренда (e.g., "Audi Q7 2023", "Audi A6 2022").
negative – автомобиль другого бренда (e.g., "BMW X5 2023").
Бренды: Audi, EXEED, Geely, Changan, Kia.
Обучение модели
Оптимизатор: AdamW (learning rate = 2e-5).
Цикл обучения: 20 эпох.
Оценка:
Матрица косинусной схожести эмбеддингов.
Метрика разделения брендов (intra_similarity vs inter_similarity).
Тестирование
Тестовые примеры: 12 автомобилей из 5 брендов.
Визуализация: Сравнение эмбеддингов до/после обучения.
Терминология:
Эмбеддинг (embedding): Векторное представление текста/объекта.
Triplet Loss: Функция потерь, использующая три элемента: anchor, positive (похож на anchor), negative (не похож).
Косинусное сходство (cosine similarity): Мера близости векторов (-1 до 1).
PCA/t-SNE: Методы снижения размерности для визуализации.
source:
```
pip install torch torchvision torchaudio scikit-learn
pip install matplotlib scikit-learn plotly umap-learn umap
pip install nomic torch transformers sentence-transformers transformers torch nomic einops
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
from sentence_transformers import SentenceTransformer
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import random
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
import plotly.express as px
import torch.nn.functional as F
def visualize_embeddings(embeddings, labels, title="Embeddings Visualization"):
"""Визуализация эмбеддингов с использованием PCA и t-SNE"""
# Уменьшение размерности
methods = {
"PCA": PCA(n_components=3),
"t-SNE": TSNE(n_components=3, perplexity=min(5, len(embeddings)-1), random_state=42),
}
# Создание фигур
num_methods = len(methods)
fig_2d, axs_2d = plt.subplots(1, num_methods, figsize=(8*num_methods, 6))
fig_3d = plt.figure(figsize=(6*num_methods, 6))
for i, (name, reducer) in enumerate(methods.items()):
# Применение метода снижения размерности
reduced = reducer.fit_transform(embeddings)
# 2D визуализация (первые два измерения)
ax = axs_2d[i]
ax.set_title(f"{name} 2D")
ax.scatter(reduced[:, 0], reduced[:, 1])
for j, label in enumerate(labels):
ax.annotate(label, (reduced[j, 0], reduced[j, 1]), fontsize=9)
# 3D визуализация
ax_3d = fig_3d.add_subplot(1, num_methods, i+1, projection='3d')
ax_3d.set_title(f"{name} 3D")
ax_3d.scatter(reduced[:, 0], reduced[:, 1], reduced[:, 2])
for j, label in enumerate(labels):
ax_3d.text(reduced[j, 0], reduced[j, 1], reduced[j, 2], label, fontsize=8)
fig_2d.suptitle(title, fontsize=16)
plt.tight_layout()
# Интерактивная 3D визуализация с Plotly
for name, reducer in methods.items():
reduced = reducer.fit_transform(embeddings)
fig = px.scatter_3d(
x=reduced[:, 0], y=reduced[:, 1], z=reduced[:, 2],
text=labels,
title=f"{title} - {name}",
labels={'x': 'X', 'y': 'Y', 'z': 'Z'}
)
fig.update_traces(textposition='top center', marker_size=5)
fig.show()
# Кастомная функция потерь с усиленным разделением брендов
class BrandAwareTripletLoss(nn.Module):
def __init__(self, margin=0.5, brand_penalty_coef=0.3):
super().__init__()
self.margin = margin
self.brand_penalty_coef = brand_penalty_coef
# Больший акцент на сохранение внутрибрендового сходства
self.intra_coef = 0.01
# Меньший, но все еще значительный акцент на разделение брендов
self.inter_coef = 0.9
def forward(self, anchor, positive, negative, anchor_texts, positives_texts, negative_texts):
# Расчет косинусного сходства
pos_sim = F.cosine_similarity(anchor, positive)
neg_sim = F.cosine_similarity(anchor, negative)
intra_loss = (1 - pos_sim).mean()
inter_loss = F.relu(neg_sim - pos_sim + self.margin).mean()
print(anchor_texts)
# print('.....')
# print(positives_texts)
# print('.....')
# print(negative_texts)
#print('.....!!!')
return self.intra_coef * intra_loss + self.inter_coef * inter_loss + (0 * 0)
# Базовый triplet loss
#loss = F.relu(neg_sim - pos_sim + self.margin)
# Дополнительный штраф за межбрендовое сходство
# brand_penalty = torch.mean(F.cosine_similarity(anchor, negative))
# Комбинированная функция потерь
#return torch.mean(loss) + self.brand_penalty_coef * brand_penalty
# Установка seed
def set_seed(seed=42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
set_seed()
# Загрузка модели
model = SentenceTransformer(
"nomic-ai/nomic-embed-text-v1.5",
trust_remote_code=True
)
# Пример данных (улучшенный набор)
triplets = [
# Базовые примеры
("Audi Q7 2023", "Audi A6 2022", "BMW X5 2023"),
("Audi Q7 2023", "Audi Q8 2023", "EXEED RX 2023"),
("Mercedes S-Class", "Mercedes E-Class", "Audi A8"),
("BMW 5 Series", "BMW 3 Series", "Lada Vesta"),
# Audi + другие бренды (усиленные примеры)
("Audi A1", "Audi A3", "EXEED LX"),
("Audi S3", "Audi RS3", "Changan Eado Plus"),
("Audi A4", "Audi S4", "Geely Emgrand"),
("Audi RS4", "Audi A5", "Kia Rio"),
("Audi e-tron", "Audi A5", "Kia EV6"),
("Audi e-tron", "Audi RS4", "Kia EV6"),
("Audi e-tron", "Audi A4", "Kia EV6"),
("Audi e-tron", "Audi Q5", "Kia EV6"),
# EXEED + другие бренды
("EXEED LX", "EXEED TXL", "Audi A1"),
("EXEED VX", "EXEED RX", "Changan CS55 Plus"),
("EXEED Sterra ES", "EXEED Sterra ET", "Geely Boyue L"),
("EXEED Sterra", "EXEED LX", "Geely Boyue L"),
("EXEED Sterra", "EXEED VX", "Geely Haoyue"),
("EXEED Yaoguang", "EXEED LX", "Kia K5"),
# Changan + другие бренды
("Changan CS75 Plus", "Changan CS85 Coupe", "Audi RS5"),
("Changan UNI-V", "Changan UNI-K", "EXEED VX"),
("Changan UNI-V", "Changan UNI-K", "Audi RS4"),
#("Changan Lumin", "Changan UNI-K", "Audi RS4"),
#("Changan BenBen", "Changan Lumin", "Audi e-tron"),
#("Changan UNI-V", "Changan Hunter", "Audi RS4"),
("Changan Raeton", "Changan Hunter", "Kia Stinger"),
# Geely + другие бренды
("Geely Tugella", "Geely Xingyue", "Audi RS7"),
("Geely Binrui", "Geely Jiaji", "EXEED Sterra ES"),
("Geely Haoyue", "Geely Geometry E", "Changan UNI-T"),
("Geely Galaxy L7", "Geely Emgrand", "Kia Seltos"),
("Geely Tugella", "Geely Monjaro", "Audi RS7"),
("Geely Tugella", "Geely Monjaro", "Changan UNI-V"),
# Kia + другие бренды
("Kia EV9", "Kia EV5", "Audi R8"),
("Kia Ceed", "Kia Rio", "EXEED Yaoguang"),
("Kia Sorento", "Kia Telluride", "Changan Alsvin"),
("Kia Niro", "Kia Soul", "Geely Geometry G6"),
("Kia Sorento", "Kia EV6", "EXEED LX"),
# ("Kia EV6", "Kia Sorento", "EXEED LX"),
# ("Kia EV6", "Kia Sorento", "Audi e-tron"),
# ("Kia EV6", "Kia EV9", "Changan Raeton"),
("Kia EV6", "Kia EV9", "Audi e-tron"),
("Audi e-tron", "Audi RS e-tron GT", "Kia EV6"),
("Audi e-tron", "Audi Q8 e-tron", "Kia EV9"),
("Audi e-tron GT", "Audi Q4 e-tron", "Kia EV6 GT"),
# Сравнение с неэлектрическими моделями
("Audi e-tron", "Audi Q5", "Kia Sportage"), # Сравнение с обычной Kia
("Kia EV6", "Kia Niro", "Audi A4"),
# Сложные случаи (похожие модели)
# ("Audi A3 Sportback", "Audi A3 Sedan", "EXEED LX SUV"),
# ("EXEED Sterra ET", "EXEED Sterra ES", "Audi e-tron"),
# ("Geely Tugella", "Geely Xingyue", "Changan CS95 Plus"),
# ("Kia EV6", "Kia EV9", "Geely Geometry A"),
# ("Kia EV6", "Kia EV9", "Audi e-tron"),
]
# Датасет
class TripletDataset(Dataset):
def __init__(self, triplets):
self.triplets = triplets
def __len__(self):
return len(self.triplets)
def __getitem__(self, idx):
return self.triplets[idx]
dataset = TripletDataset(triplets)
train_dataloader = DataLoader(dataset, shuffle=True, batch_size=5, drop_last=True )
# Устройство
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.train()
# Функция потерь и оптимизатор
triplet_loss = BrandAwareTripletLoss(margin=0.5, brand_penalty_coef=0.5)
optimizer = AdamW(model.parameters(), lr=2e-5, weight_decay=0.01) # Более низкий lr
# Тестовая функция
def test_model(model, test_texts):
model.eval()
with torch.no_grad():
embeddings = model.encode(test_texts)
sim_matrix = cosine_similarity(embeddings)
print("\n🔍 Матрица сходства:")
for i, row in enumerate(sim_matrix):
print(f"{test_texts[i][:15]:<15} " + " ".join([f"{x:.2f}" for x in row]))
# Визуализация только для финального теста
if epoch == 19:
visualize_embeddings(embeddings, test_texts, "Final Embeddings")
model.train()
return sim_matrix
# Функция для оценки разделения брендов
def evaluate_brand_separation(model, brands):
model.eval()
with torch.no_grad():
embeddings = model.encode(brands)
intra_similarity, inter_similarity = 0, 0
count_intra, count_inter = 0, 0
for i, brand1 in enumerate(brands):
for j, brand2 in enumerate(brands):
if i == j:
continue
sim = F.cosine_similarity(
torch.tensor(embeddings[i]).unsqueeze(0),
torch.tensor(embeddings[j]).unsqueeze(0)
).item()
brand1_prefix = brand1.split()[0]
brand2_prefix = brand2.split()[0]
if brand1_prefix == brand2_prefix:
intra_similarity += sim
count_intra += 1
else:
inter_similarity += sim
count_inter += 1
intra_similarity /= count_intra if count_intra > 0 else 1
inter_similarity /= count_inter if count_inter > 0 else 1
print(f"Среднее сходство внутри брендов: {intra_similarity:.4f}")
print(f"Среднее сходство между брендами: {inter_similarity:.4f}")
print(f"Разница: {intra_similarity - inter_similarity:.4f}")
model.train()
return intra_similarity - inter_similarity
# Тестовые данные
test_texts = [
"Audi A4", "Audi Q5", "Audi e-tron",
"EXEED LX", "EXEED VX", "EXEED Sterra",
"Geely Tugella", "Geely Monjaro",
"Changan CS75", "Changan UNI-K",
"Kia EV6", "Kia Sorento"
]
print("🧪 Тест ДО обучения:")
test_model(model, test_texts)
brand_diff_before = evaluate_brand_separation(model, test_texts)
# Обучение
print("\n🔁 Начало обучения...")
num_epochs = 20
for epoch in range(num_epochs):
total_loss = 0
for batch in train_dataloader:
print(batch)
anchors, positives, negatives = batch
#anchors = [item[0] for item in batch]
#positives = [item[1] for item in batch]
#negatives = [item[2] for item in batch]
print(f"Батч , размер: {len(batch)} триплетов")
print(f"Anchors: {len(anchors)} элементов")
print(f"Positives: {len(positives)} элементов")
print(f"Negatives: {len(negatives)} элементов")
print('!!!!')
all_texts = anchors + positives + negatives
features = model.tokenize(all_texts)
features = {k: v.to(device) for k, v in features.items()}
output = model(features)
all_embeddings = output['sentence_embedding']
all_embeddings = F.normalize(all_embeddings, p=2, dim=1)
batch_size = len(anchors)
anchor_emb = all_embeddings[:batch_size]
positive_emb = all_embeddings[batch_size:2*batch_size]
negative_emb = all_embeddings[2*batch_size:]
# Расчет потерь с использованием кастомной функции
loss = triplet_loss(anchor_emb, positive_emb, negative_emb, anchors, positives, negatives)
total_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
avg_loss = total_loss / len(train_dataloader)
print(f"Эпоха {epoch+1}/{num_epochs} | Потери: {avg_loss:.4f}")
# Промежуточная оценка
if True: #(epoch + 1) % 5 == 0:
print(f"\n🧪 Промежуточный тест после эпохи {epoch+1}:")
#test_model(model, test_texts)
brand_diff = evaluate_brand_separation(model, test_texts)
print(f"Улучшение разделения брендов: {brand_diff - brand_diff_before:.4f}")
print("\n🧪 Финальный тест ПОСЛЕ обучения:")
test_model(model, test_texts)
brand_diff_after = evaluate_brand_separation(model, test_texts)
print(f"Общее улучшение разделения брендов: {brand_diff_after - brand_diff_before:.4f}")
# Сохранение модели
output_path = "./brand_aware_nomic_embed"
model.save(output_path)
print(f"\n✅ Модель сохранена в {output_path}")
```
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment