Update app.py

app.py

@@ -1,14 +1,29 @@
+import torch
+import numpy as np
+from torch.nn.utils.rnn import pad_sequence
 import gradio as gr
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 # Load the model and tokenizer
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model_name = "google/flan-t5-large"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = SentenceTransformer('AnnaWegmann/Style-Embedding', device='cpu')
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+model.to(device)
 
-def run_tinystyler(source_text, target_example_texts, reranking, temperature, top_p):
-    concatenated_text = source_text + " " + target_example_texts
-    inputs = tokenizer(concatenated_text, return_tensors="pt")
+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]
+    embeddings = model.encode(all_target_texts, batch_size=len(all_target_texts), convert_to_tensor=True, show_progress_bar=False)
+    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)
+    mean_embeddings = torch.sum(padded_embeddings * mask, dim=1) / mask.sum(dim=1)
+    return mean_embeddings.cpu().numpy()
+
+def run_tinystyler_batch(source_texts, target_example_texts_batch, reranking, temperature, top_p):
+    inputs = tokenizer(source_texts, return_tensors="pt")
     
     # Generate the output with specified temperature and top_p
     output = model.generate(
@@ -19,8 +34,11 @@ def run_tinystyler(source_text, target_example_texts, reranking, temperature, to
         max_length=1024
     )
     
-    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return generated_text
+    generated_texts = tokenizer.decode_batch(output, skip_special_tokens=True)
+    return generated_texts
+    
+def run_tinystyler(source_text, target_example_texts, reranking, temperature, top_p):
+    return run_tinystyler_batch([source_text], [target_example_texts], reranking, temperature, top_p)[0]
 
 # Preset examples with cached generations
 preset_examples = {