Update classification.py

classification.py

@@ -173,12 +173,12 @@ def process_categories(categories, model):
 def match_categories(df, category_df, treshold=0.45):
     for topic in category_df['topic']:
         df[topic] = 0
-    for i, ebd_content in enumerate(df['Embeddings']):
+    for index, ebd_content in enumerate(df['Embeddings']):
         if isinstance(ebd_content, torch.Tensor):
             cos_scores = util.cos_sim(ebd_content, torch.stack(list(category_df['Embeddings']), dim=0))[0]
             high_score_indices = [i for i, score in enumerate(cos_scores) if score > treshold]
             for j in high_score_indices:
-                df.loc[i, category_df.loc[j, 'topic']] = float(cos_scores[j])
+                df.loc[index, category_df.loc[j, 'topic']] = 'float(cos_scores[j])'
     return df
 
 def save_data(df, filename):
@@ -193,9 +193,10 @@ def classification(column, file_path, categories, treshold):
 
     # Initialize models
     model_ST = initialize_models()
-
+    print('Generating Embeddings')
     # Generate embeddings for df
     df = generate_embeddings(df, model_ST, column)
+    print('Embeddings Generated')
 
 
     category_df = process_categories(categories, model_ST)