Buckets:
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from typing import List, Dict, Tuple, Any, Optional | |
| class GeometricProblemEncoder(nn.Module): | |
| def __init__(self, vocab_size=5000, embedding_dim=256, hidden_dim=512): | |
| super().__init__() | |
| self.text_embedding = nn.Embedding(vocab_size, embedding_dim) | |
| self.text_encoder = nn.LSTM( | |
| embedding_dim, | |
| hidden_dim, | |
| num_layers=2, | |
| batch_first=True, | |
| bidirectional=True, | |
| ) | |
| self.object_types = ["point", "line", "circle", "triangle", "angle"] | |
| self.object_embedding = nn.Embedding(len(self.object_types), embedding_dim) | |
| self.diagram_encoder = nn.Sequential( | |
| nn.Conv2d(3, 64, kernel_size=3, padding=1), | |
| nn.ReLU(), | |
| nn.MaxPool2d(2), | |
| nn.Conv2d(64, 128, kernel_size=3, padding=1), | |
| nn.ReLU(), | |
| nn.MaxPool2d(2), | |
| nn.Conv2d(128, 256, kernel_size=3, padding=1), | |
| nn.ReLU(), | |
| nn.AdaptiveAvgPool2d((4, 4)), | |
| nn.Flatten(), | |
| nn.Linear(256 * 16, hidden_dim), | |
| ) | |
| self.fusion = nn.Sequential( | |
| nn.Linear(hidden_dim * 3, hidden_dim), | |
| nn.ReLU(), | |
| nn.Dropout(0.1), | |
| nn.Linear(hidden_dim, hidden_dim), | |
| ) | |
| def encode_text(self, text_tokens): | |
| embedded = self.text_embedding(text_tokens) | |
| output, (hidden, _) = self.text_encoder(embedded) | |
| final_hidden = torch.cat([hidden[-2], hidden[-1]], dim=-1) | |
| return final_hidden | |
| def encode_objects(self, object_types): | |
| indices = [] | |
| for obj_type in object_types: | |
| if obj_type in self.object_types: | |
| indices.append(self.object_types.index(obj_type)) | |
| if not indices: | |
| return torch.zeros(1, self.text_encoder.hidden_size * 2) | |
| indices_tensor = torch.LongTensor(indices) | |
| embedded = self.object_embedding(indices_tensor) | |
| return embedded.mean(dim=0, keepdim=True) | |
| def encode_diagram(self, diagram_image): | |
| if diagram_image is None: | |
| return torch.zeros(1, self.text_encoder.hidden_size * 2) | |
| encoded = self.diagram_encoder(diagram_image) | |
| return encoded | |
| def forward(self, text_tokens, object_types, diagram_image=None): | |
| text_repr = self.encode_text(text_tokens) | |
| object_repr = self.encode_objects(object_types) | |
| diagram_repr = self.encode_diagram(diagram_image) | |
| batch_size = text_repr.size(0) | |
| if object_repr.size(0) != batch_size: | |
| object_repr = object_repr.expand(batch_size, -1) | |
| if diagram_repr.size(0) != batch_size: | |
| diagram_repr = diagram_repr.expand(batch_size, -1) | |
| combined = torch.cat([text_repr, object_repr, diagram_repr], dim=-1) | |
| fused_repr = self.fusion(combined) | |
| return fused_repr | |
Xet Storage Details
- Size:
- 2.95 kB
- Xet hash:
- 45c9f2596dafc6d2ab03a21f6f4a854b140437318b4ae75442890c500077a9aa
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.