Update main.py

main.py

@@ -105,9 +105,9 @@ class UnifiedModel(nn.Module):
                 input_ids=input_id,
                 attention_mask=attn_mask
             )
-            hidden_states.append(outputs.pooler_output) # Usar pooler_output para obtener un vector de 768
+            hidden_states.append(outputs.logits[:, -1, :])  # Obtener el último hidden state
 
-        concatenated_hidden_states = torch.cat(hidden_states, dim=1) # Concatenar en la dimensión correcta
+        concatenated_hidden_states = torch.cat(hidden_states, dim=1)
         logits = self.classifier(concatenated_hidden_states)
         return logits
 
@@ -390,7 +390,8 @@ async def get_home():
 def push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name):
     for model_name, model in models.items():
         torch.save(model.state_dict(), model_name)
-        redis_client.set(f"model:{model_name}", open(model_name, "rb").read())
+        with open(model_name, "rb") as f:
+            redis_client.set(f"model:{model_name}", f.read())
 
     for tokenizer_name, tokenizer in tokenizers.items():
         tokens = tokenizer.get_vocab()