Spaces:
Paused
Paused
File size: 5,209 Bytes
9189e38 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
import re
import csv
import random
import numpy as np
import torch
from tqdm import tqdm
from transformers import pipeline
from sentence_transformers import SentenceTransformer, util
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
from db.db_utils import get_connection, get_mapping_from_db
# Load a pre-trained SBERT model
model = SentenceTransformer('all-MiniLM-L6-v2')
# Set seeds for reproducibility of zero-shot classification
def set_seed(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
set_seed(1)
# Load a pre-trained model and tokenizer
classifier = pipeline("zero-shot-classification", model="roberta-large-mnli")
# Load food categories from CSV
def load_food_categories(csv_file):
food_categories = set()
with open(csv_file, newline='') as csvfile:
reader = csv.DictReader(csvfile)
for row in reader:
food_categories.add(row['food_category'])
return list(food_categories)
# Path to the CSV file with food categories
csv_file_path = 'dictionary/dictionary.csv'
food_categories = load_food_categories(csv_file_path)
# Precompute embeddings for food categories
category_embeddings = model.encode(food_categories, convert_to_tensor=True)
# Classify item as food or non-food
def classify_item(item, cursor):
# Check database for item
cleaned_item = item.strip().lower()
mapping = get_mapping_from_db(cursor, cleaned_item)
if mapping and 'is_food' in mapping:
is_food = mapping['is_food']
if is_food is not None:
print(f"Item: {item} found in database with is_food: {is_food}")
return ("food" if is_food else "non-food"), 1.0
# If not found in database, classify using the model
result = classifier(item, candidate_labels=["food", "non-food"])
label = result["labels"][0]
score = result["scores"][0]
print(f"Item: {item}, Label: {label}, Score: {score}")
return label, score
# Determine the category of a food item
def determine_category(item):
item_embedding = model.encode(item, convert_to_tensor=True)
similarities = util.pytorch_cos_sim(item_embedding, category_embeddings)
category_idx = similarities.argmax()
category = food_categories[category_idx]
# Assuming 'similarities' is a tensor of similarity scores and 'food_categories' is the list of category names
top_3_indices = torch.topk(similarities, 3).indices[0].tolist()
top_3_scores = torch.topk(similarities, 3).values[0].tolist()
top_3_categories = [(food_categories[idx], score) for idx, score in zip(top_3_indices, top_3_scores)]
print("=========================================")
print(f"item: {item}")
for category, score in top_3_categories:
print(f"Category: {category}, Similarity Score: {score:.4f}")
return category
# Visualize embeddings
def visualize_embeddings(items, categories, item_embeddings, category_embeddings):
tsne = TSNE(n_components=2, random_state=1)
embeddings = torch.cat([item_embeddings, category_embeddings], dim=0)
tsne_embeddings = tsne.fit_transform(embeddings.detach().cpu().numpy())
plt.figure(figsize=(10, 10))
for i, label in enumerate(items + categories):
x, y = tsne_embeddings[i]
plt.scatter(x, y)
plt.text(x+0.1, y+0.1, label, fontsize=9)
plt.show()
# Categorize food items
def categorize_food_items(items):
categories_found = set()
for item in items:
category = determine_category(item)
categories_found.add(category)
print(f"Categories found: {categories_found}")
if len(categories_found) == 1:
return list(categories_found)[0]
elif len(categories_found) > 1:
return "heterogeneous mixture"
else:
return "food"
# List of items to test
items_to_test = [
"Swiss Cheese, Provolone cheese, cheddar, mozzarella"
]
# Initialize database connection
conn = get_connection()
cursor = conn.cursor()
# Collect results and visualize embeddings
results = []
item_embeddings = []
items_list = []
for items in items_to_test:
# Split items by both "/" and "," and strip extra spaces
items_list = [item.strip().lower() for item in re.split(r'[\/,]', items)]
item_labels = [classify_item(item, cursor) for item in items_list]
non_food_items = [item for item, (label, _) in zip(items_list, item_labels) if label == "non-food"]
# Get embeddings for visualization
item_embeddings.extend(model.encode(items_list, convert_to_tensor=True))
list_label = categorize_food_items(items_list)
results.append([items, list_label, ", ".join(non_food_items)])
# Visualize embeddings
visualize_embeddings(items_list, food_categories, torch.stack(item_embeddings), category_embeddings)
# Write results to a CSV file
with open('multi-item-experiments/classification_results2.csv', 'w', newline='') as csvfile:
csvwriter = csv.writer(csvfile)
csvwriter.writerow(['Item List', 'Category', 'Non-Food Items'])
csvwriter.writerows(results)
# Close the SQLite connection
conn.close()
|