Update app.py

app.py

@@ -28,12 +28,14 @@ st.set_page_config(layout="wide")
 
 classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
 
-
+BATCH_SIZE = 100
 #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()  
+    probabilities = []
+    for i in range(0, len(reviews), BATCH_SIZE):
+        inputs = tokenizer(reviews[i:i+BATCH_SIZE], return_tensors='pt', truncation=True, padding=True, max_length=512)
+        outputs = model(**inputs)
+        probabilities.extend(F.softmax(outputs.logits, dim=1).tolist())
     return probabilities
 
 def top_rating(scores):
@@ -65,12 +67,14 @@ def process_filter_words(filter_words_input):
 def classify_with_new_classes(reviews, class_names):
     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
 
@@ -139,14 +143,9 @@ def process_reviews(df, review_column, class_names):
         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)
-            raw_scores.extend(batch_scores)
-            review_counter += len(batch_reviews)
+        raw_scores = classify_reviews(df[review_column].tolist())
+        for i in range(0, len(raw_scores), BATCH_SIZE):
+            review_counter += BATCH_SIZE
             progress_bar.progress(review_counter / total_reviews)
 
     with st.spinner('Generating classes...'):