Update app.py

app.py

@@ -234,45 +234,31 @@ def classify_emotion(text, classifier):
     return final_emotion
         
 def get_embedding_for_text(text, tokenizer, model):
-    """Get embedding for complete text by processing in chunks."""
-    # Tokenize the entire text
-    tokens = tokenizer.tokenize(text)
-    
-    # Process in chunks of 510 tokens (512 - 2 special tokens)
-    chunk_size = 510
+    """Get embedding for complete text."""
+    chunks = split_text(text)
     chunk_embeddings = []
     
-    for i in range(0, len(tokens), chunk_size):
-        # Take a chunk of tokens
-        chunk_tokens = tokens[i:i + chunk_size]
-        
-        # Add special tokens
-        chunk_tokens = ['[CLS]'] + chunk_tokens + ['[SEP]']
-        
-        # Convert to input IDs
-        input_ids = tokenizer.convert_tokens_to_ids(chunk_tokens)
-        
-        # Pad to 512 tokens
-        input_ids += [tokenizer.pad_token_id] * (512 - len(input_ids))
-        
-        # Create attention mask
-        attention_mask = [1] * len(chunk_tokens) + [0] * (512 - len(chunk_tokens))
-        
-        # Convert to tensors
-        input_ids = torch.tensor([input_ids])
-        attention_mask = torch.tensor([attention_mask])
+    for chunk in chunks:
+        inputs = tokenizer(
+            chunk,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=512
+        )
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
         
-        # Get embedding for this chunk
         with torch.no_grad():
-            outputs = model(input_ids, attention_mask=attention_mask)
-            chunk_embedding = outputs[0][:, 0, :].cpu().numpy()
-            chunk_embeddings.append(chunk_embedding[0])
+            # Access the first element of the tuple directly
+            outputs = model(**inputs)
+            embedding = outputs[0][:, 0, :].cpu().numpy()
+            chunk_embeddings.append(embedding[0])
     
-    # Average embeddings from all chunks
     if chunk_embeddings:
-        return np.mean(chunk_embeddings, axis=0)
-    
-    # Fallback if no embeddings could be generated
+        weights = np.array([len(chunk.split()) for chunk in chunks])
+        weights = weights / weights.sum()
+        weighted_embedding = np.average(chunk_embeddings, axis=0, weights=weights)
+        return weighted_embedding
     return np.zeros(model.config.hidden_size)
     
 def format_topics(topic_model, topic_counts):