Update app.py

app.py

@@ -30,12 +30,19 @@ classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnl
 
 
 #defs
-def classify_reviews(reviews):
-    inputs = tokenizer(reviews, return_tensors='pt', truncation=True, padding=True, max_length=512)
-    outputs = model(**inputs)
-    probabilities = F.softmax(outputs.logits, dim=1).tolist()  
+def classify_reviews(reviews, batch_size=64):
+    probabilities = []
+    
+    for i in range(0, len(reviews), batch_size):
+        batch_reviews = reviews[i:i+batch_size]
+        inputs = tokenizer(batch_reviews, return_tensors='pt', truncation=True, padding=True, max_length=512)
+        outputs = model(**inputs)
+        batch_probabilities = F.softmax(outputs.logits, dim=1).tolist()
+        probabilities.extend(batch_probabilities)
+        
     return probabilities
 
+
 def top_rating(scores):
     return scores.index(max(scores)) + 1  
 
@@ -62,15 +69,17 @@ def process_filter_words(filter_words_input):
 
 
 # Function for classifying with the new model
-def classify_with_new_classes(reviews, class_names):
+def classify_with_new_classes(reviews, class_names, batch_size=64):
     class_scores = []
 
-    for review in reviews:
-        result = classifier(review, class_names)
-        scores_dict = dict(zip(result['labels'], result['scores']))
-        # Reorder scores to match the original class_names order
-        scores = [scores_dict[name] for name in class_names]
-        class_scores.append(scores)
+    for i in range(0, len(reviews), batch_size):
+        batch_reviews = reviews[i:i+batch_size]
+        for review in batch_reviews:
+            result = classifier(review, class_names)
+            scores_dict = dict(zip(result['labels'], result['scores']))
+            # Reorder scores to match the original class_names order
+            scores = [scores_dict[name] for name in class_names]
+            class_scores.append(scores)
 
     return class_scores
 
@@ -133,24 +142,27 @@ def main():
 
 
 
-def process_reviews(df, review_column, class_names):
+def process_reviews(df, review_column, class_names, batch_size=64):
     with st.spinner('Classifying reviews...'):
         progress_bar = st.progress(0)
         total_reviews = len(df[review_column].tolist())
         review_counter = 0
 
-        batch_size = 50
         raw_scores = []
         reviews = df[review_column].tolist()
         for i in range(0, len(reviews), batch_size):
             batch_reviews = reviews[i:i+batch_size]
-            batch_scores = classify_reviews(batch_reviews)
+            batch_scores = classify_reviews(batch_reviews, batch_size)
             raw_scores.extend(batch_scores)
             review_counter += len(batch_reviews)
             progress_bar.progress(review_counter / total_reviews)
 
     with st.spinner('Generating classes...'):
-        class_scores = classify_with_new_classes(df[review_column].tolist(), class_names)
+        class_scores = []
+        for i in range(0, len(reviews), batch_size):
+            batch_reviews = reviews[i:i+batch_size]
+            batch_scores = classify_with_new_classes(batch_reviews, class_names, batch_size)
+            class_scores.extend(batch_scores)
     
     class_scores_dict = {}  # New dictionary to store class scores
     for i, name in enumerate(class_names):
@@ -161,7 +173,6 @@ def process_reviews(df, review_column, class_names):
     if class_names and not all(name.isspace() for name in class_names):
         df['Highest Class'] = df[class_names].idxmax(axis=1)
 
-
     df_new = df.copy()
     df_new['raw_scores'] = raw_scores
     scores_to_df(df_new)
@@ -181,6 +192,7 @@ def process_reviews(df, review_column, class_names):
 
 
 
+
 def scores_to_df(df):
     for i in range(1, 6):
         df[f'{i} Star'] = df['raw_scores'].apply(lambda scores: scores[i-1]).round(2)