Update app.py

app.py

@@ -9,9 +9,12 @@ from sentence_transformers import SentenceTransformer
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model_name = "google/flan-t5-large"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-embedding_model = SentenceTransformer('AnnaWegmann/Style-Embedding', device='cpu')
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 model.to(device)
+embedding_model = SentenceTransformer('AnnaWegmann/Style-Embedding', device='cpu').half()
+luar_model = AutoModel.from_pretrained("rrivera1849/LUAR-MUD", trust_remote_code=True).half()
+luar_tokenizer = AutoTokenizer.from_pretrained("rrivera1849/LUAR-MUD", trust_remote_code=True)
+
 
 def get_target_style_embeddings(target_texts_batch):
     all_target_texts = [target_text for target_texts in target_texts_batch for target_text in target_texts]
@@ -19,13 +22,29 @@ def get_target_style_embeddings(target_texts_batch):
     lengths = [len(target_texts) for target_texts in target_texts_batch]
     split_embeddings = torch.split(embeddings, lengths)
     padded_embeddings = pad_sequence(split_embeddings, batch_first=True, padding_value=0.0)
-    mask = (torch.arange(padded_embeddings.size(1))[None, :] < torch.tensor(lengths)[:, None]).to(torch.float32).unsqueeze(-1)
+    mask = (torch.arange(padded_embeddings.size(1))[None, :] < torch.tensor(lengths)[:, None]).to(embeddings.dtype).unsqueeze(-1)
     mean_embeddings = torch.sum(padded_embeddings * mask, dim=1) / mask.sum(dim=1)
-    return mean_embeddings.cpu().numpy()
+    return mean_embeddings.float().cpu().numpy()
+
+def get_luar_embeddings(texts_batch):
+    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]
+    max_episode_length = max(episode_lengths)
+    sequence_lengths = [t["attention_mask"].shape[1] for t in tokenized_episodes]
+    max_sequence_length = max(sequence_lengths)
+    padded_input_ids = [torch.nn.functional.pad(t["input_ids"], (0, 0, 0, max_episode_length - t["input_ids"].shape[0])) for t in tokenized_episodes]
+    padded_attention_mask = [torch.nn.functional.pad(t["attention_mask"], (0, 0, 0, max_episode_length - t["attention_mask"].shape[0])) for t in tokenized_episodes]
+    input_ids = torch.stack(padded_input_ids)
+    attention_mask = torch.stack(padded_attention_mask)
+    return luar_model(input_ids=input_ids, attention_mask=attention_mask).float().cpu().numpy()
 
 def run_tinystyler_batch(source_texts, target_texts_batch, reranking, temperature, top_p):
-    inputs = tokenizer(source_texts, return_tensors="pt")
+    inputs = tokenizer(source_texts, return_tensors="pt").to(device)
     target_style_embeddings = get_target_style_embeddings(target_texts_batch)
+    source_style_luar_embeddings = get_luar_embeddings([[st] for st in source_texts])
+    target_style_luar_embeddings = get_luar_embeddings(target_texts_batch)
+    
     
     # Generate the output with specified temperature and top_p
     output = model.generate(