Update app.py

app.py

@@ -2,7 +2,7 @@ import torch
 import numpy as np
 from torch.nn.utils.rnn import pad_sequence
 import gradio as gr
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import AutoModel, AutoModelForSeq2SeqLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
 
 # Load the model and tokenizer
@@ -28,6 +28,7 @@ def get_target_style_embeddings(target_texts_batch):
     return mean_embeddings.float().cpu().numpy()
 
 def get_luar_embeddings(texts_batch):
+    assert set([len(texts) for texts in texts_batch]) == 1
     episodes = texts_batch
     tokenized_episodes = [luar_tokenizer(episode, max_length=512, padding="longest", truncation=True, return_tensors="pt").to(device) for episode in episodes]
     episode_lengths = [t["attention_mask"].shape[0] for t in tokenized_episodes]