Update classification.py

classification.py

@@ -17,9 +17,15 @@ def initialize_models():
 def generate_embeddings(df, model, Column):
     embeddings_list = []
     for index, row in df.iterrows():
-        if type(row["Title"]) == str and type(row[Column]) == str:
+        if type(row[Column]) == str:
             print(index)
-            content = row["Title"] + "\n" + row[Column]
+            if 'Title' in df.columns:
+                if type(row["Title"]) == str:
+                    content = row["Title"] + "\n" + row[Column]
+                else:
+                    content = row[Column]
+            else:
+                content = row[Column]
             embeddings = model.encode(content, convert_to_tensor=True)
             embeddings_list.append(embeddings)
         else:
@@ -39,13 +45,13 @@ def process_categories(categories, model):
 
 
 
-def match_categories(df, category_df):
+def match_categories(df, category_df, treshold=0.45):
     
     categories_list, experts_list, topic_list, scores_list = [], [], [], []
     for ebd_content in 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 > 0.45]
+            high_score_indices = [i for i, score in enumerate(cos_scores) if score > treshold]
 
             # Append the corresponding categories, experts, and topics for each high-scoring index
             categories_list.append([category_df.loc[index, 'description'] for index in high_score_indices])
@@ -86,7 +92,7 @@ def save_data(df, filename):
     df.to_excel(new_filename, index=False)
     return new_filename
 
-def classification(column, file_path, categories):
+def classification(column, file_path, categories, treshold):
     # Load data
     df = load_data(file_path)
 
@@ -100,7 +106,7 @@ def classification(column, file_path, categories):
     category_df = process_categories(categories, model_ST)
 
     # Match categories
-    df = match_categories(df, category_df)
+    df = match_categories(df, category_df, treshold=treshold)
 
     # Save data
     return save_data(df,file_path), df