shape

app.py

@@ -36,11 +36,12 @@ def predict(type, normalized_asm):
     
     tokenizer_output = nova_tokenizer.encode(inputs, '', char_types)
     input_ids = torch.LongTensor(tokenizer_output['input_ids'].tolist()).unsqueeze(0)
+    print("Input IDs:", input_ids.shape)
     nova_attention_mask = torch.LongTensor(tokenizer_output['nova_attention_mask']).unsqueeze(0)
 
     output = model.generate(
         inputs=input_ids.cuda(), max_new_tokens=512, temperature=0.2, top_p=0.95,
-        num_return_sequences=20, do_sample=True, nova_attention_mask=nova_attention_mask.cuda(),
+        num_return_sequences=1, do_sample=True, nova_attention_mask=nova_attention_mask.cuda(),
         no_mask_idx=torch.LongTensor([tokenizer_output['no_mask_idx']]).cuda(), 
         pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id
     )