Update app.py

app.py

@@ -221,24 +221,14 @@ def generate_30(
     pooled_prompt_embeds_b = prompt_embeds_b[0]  # Pooled output from encoder 2
     prompt_embeds_b = prompt_embeds_b.hidden_states[-2]  # Penultimate hidden state from encoder 2
 
-    # 3. Concatenate the embeddings along the sequence dimension (dim=1)
-    prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b], dim=1)
+    # 3. Concatenate the embeddings 
+    prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b])
     print('catted shape: ', prompt_embeds.shape)
-    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=1)
-    
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b])
     # 4. (Optional) Average the pooled embeddings
-    prompt_embeds = prompt_embeds.mean(dim=0)
+    prompt_embeds = torch.mean(prompt_embeds,dim=0,keepdim=True)
     print('averaged shape: ', prompt_embeds.shape)
-    test_prompt_embeds = prompt_embeds.mean(dim=0,keepdim=True)
-    print('averaged shape (keepdim): ', test_prompt_embeds.shape)
-
-    test_prompt_embeds_2 = torch.cat([prompt_embeds_a, prompt_embeds_b], dim=0).mean(dim=1)
-    print('averaged shape 2: ', test_prompt_embeds_2.shape)
-    test_prompt_embeds_3 = torch.cat([prompt_embeds_a, prompt_embeds_b]).mean(dim=0,keepdim=True)
-    print('averaged shape 3(keepdim): ', test_prompt_embeds_3.shape)
-    
-    pooled_prompt_embeds = pooled_prompt_embeds.mean(dim=0)
-
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0,keepdim=True)
 
     options = {
         #"prompt": prompt,
@@ -314,23 +304,14 @@ def generate_60(
     pooled_prompt_embeds_b = prompt_embeds_b[0]  # Pooled output from encoder 2
     prompt_embeds_b = prompt_embeds_b.hidden_states[-2]  # Penultimate hidden state from encoder 2
 
-    # 3. Concatenate the embeddings along the sequence dimension (dim=1)
-    prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b], dim=1)
+    # 3. Concatenate the embeddings 
+    prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b])
     print('catted shape: ', prompt_embeds.shape)
-    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=1)
-    
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b])
     # 4. (Optional) Average the pooled embeddings
-    prompt_embeds = prompt_embeds.mean(dim=0)
+    prompt_embeds = torch.mean(prompt_embeds,dim=0,keepdim=True)
     print('averaged shape: ', prompt_embeds.shape)
-    test_prompt_embeds = prompt_embeds.mean(dim=0,keepdim=True)
-    print('averaged shape (keepdim): ', prompt_embeds.shape)
-
-    test_prompt_embeds_2 = torch.cat([prompt_embeds_a, prompt_embeds_b], dim=1).mean(dim=1)
-    print('averaged shape 2: ', test_prompt_embeds_2.shape)
-    test_prompt_embeds_3 = torch.cat([prompt_embeds_a, prompt_embeds_b]).mean(dim=1,keepdim=True)
-    print('averaged shape 3(keepdim): ', test_prompt_embeds_3.shape)
-    
-    pooled_prompt_embeds = pooled_prompt_embeds.mean(dim=0)
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0,keepdim=True)
 
 
     options = {
@@ -407,24 +388,14 @@ def generate_90(
     pooled_prompt_embeds_b = prompt_embeds_b[0]  # Pooled output from encoder 2
     prompt_embeds_b = prompt_embeds_b.hidden_states[-2]  # Penultimate hidden state from encoder 2
 
-    # 3. Concatenate the embeddings along the sequence dimension (dim=1)
-    prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b], dim=1)
+    # 3. Concatenate the embeddings 
+    prompt_embeds = torch.cat([prompt_embeds_a, prompt_embeds_b])
     print('catted shape: ', prompt_embeds.shape)
-    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=1)
-    
+    pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_a, pooled_prompt_embeds_b])
     # 4. (Optional) Average the pooled embeddings
-    prompt_embeds = prompt_embeds.mean(dim=0)
+    prompt_embeds = torch.mean(prompt_embeds,dim=0,keepdim=True)
     print('averaged shape: ', prompt_embeds.shape)
-    test_prompt_embeds = prompt_embeds.mean(dim=0,keepdim=True)
-    print('averaged shape (keepdim): ', prompt_embeds.shape)
-
-    test_prompt_embeds_2 = torch.cat([prompt_embeds_a, prompt_embeds_b], dim=1).mean(dim=1)
-    print('averaged shape 2: ', test_prompt_embeds_2.shape)
-    test_prompt_embeds_3 = torch.cat([prompt_embeds_a, prompt_embeds_b]).mean(dim=1,keepdim=True)
-    print('averaged shape 3(keepdim): ', test_prompt_embeds_3.shape)
-    
-    pooled_prompt_embeds = pooled_prompt_embeds.mean(dim=0)
-
+    pooled_prompt_embeds = torch.mean(pooled_prompt_embeds,dim=0,keepdim=True)
 
     options = {
         #"prompt": prompt,