Update app.py

app.py

@@ -4,6 +4,7 @@ from torch.nn.utils.rnn import pad_sequence
 import gradio as gr
 from transformers import AutoModel, AutoModelForSeq2SeqLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
+from time import time
 
 # Load the model and tokenizer
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -46,7 +47,9 @@ def run_tinystyler_batch(source_texts, target_texts_batch, reranking, temperatur
     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])
+    print("Log 0", time(), source_style_luar_embeddings.shape)
     target_style_luar_embeddings = get_luar_embeddings(target_texts_batch)
+    print("Log 1", time(), target_style_luar_embeddings.shape)
     
     
     # Generate the output with specified temperature and top_p
@@ -58,14 +61,16 @@ def run_tinystyler_batch(source_texts, target_texts_batch, reranking, temperatur
         max_length=1024,
         num_return_sequences=reranking,
     )
-    
+    print("Log 2", time(), output.shape)
     generated_texts = tokenizer.batch_decode(output, skip_special_tokens=True)
 
     # Evaluate candidates
     candidates_luar_embeddings = [get_luar_embeddings([[candidates[i]] for candidates in generated_texts]) for i in range(reranking)]
+    print("Log 3", time(), len(candidates_luar_embeddings), len(candidates_luar_embeddings[0]))
 
     # Get best based on re-ranking
     generated_texts = [texts[0] for texts in generated_texts]
+    print("Final Log", time(), len(generated_texts))
     
     return generated_texts