tanner/min-dalle-test
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

fixed typing error for older python versions

Browse Source
This commit is contained in:
Brett Kuprel
2022-07-02 09:06:22 -04:00
parent 2dadfdfb31
commit 313635e914
5 changed files with 16 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
min_dalle/models/dalle_bart_decoder.py
View File

@@ -1,3 +1,4 @@
from typing import Tuple, List
import torch
from torch import LongTensor, nn, FloatTensor, BoolTensor
torch.set_grad_enabled(False)
@@ -25,7 +26,7 @@ class DecoderSelfAttention(AttentionBase):
attention_state: FloatTensor,
attention_mask: BoolTensor,
token_mask: BoolTensor
) -> tuple[FloatTensor, FloatTensor]:
) -> Tuple[FloatTensor, FloatTensor]:
keys = self.k_proj.forward(decoder_state)
values = self.v_proj.forward(decoder_state)
queries = self.q_proj.forward(decoder_state)
@@ -70,7 +71,7 @@ class DecoderLayer(nn.Module):
attention_state: FloatTensor,
attention_mask: BoolTensor,
token_index: LongTensor
) -> tuple[FloatTensor, FloatTensor]:
) -> Tuple[FloatTensor, FloatTensor]:
# Self Attention
residual = decoder_state
decoder_state = self.pre_self_attn_layer_norm.forward(decoder_state)
@@ -125,7 +126,7 @@ class DalleBartDecoder(nn.Module):
self.image_token_count = image_token_count
self.embed_tokens = nn.Embedding(image_vocab_count + 1, embed_count)
self.embed_positions = nn.Embedding(image_token_count, embed_count)
self.layers: list[DecoderLayer] = nn.ModuleList([
self.layers: List[DecoderLayer] = nn.ModuleList([
DecoderLayer(
image_token_count,
attention_head_count,
@@ -153,7 +154,7 @@ class DalleBartDecoder(nn.Module):
attention_state: FloatTensor,
prev_tokens: LongTensor,
token_index: LongTensor
) -> tuple[LongTensor, FloatTensor]:
) -> Tuple[LongTensor, FloatTensor]:
image_count = encoder_state.shape[0] // 2
token_index_batched = token_index[[0] * image_count * 2]
prev_tokens = prev_tokens[list(range(image_count)) * 2]
@@ -209,7 +210,7 @@ class DalleBartDecoder(nn.Module):
if torch.cuda.is_available(): attention_state = attention_state.cuda()
image_tokens = self.start_token[[0] * image_count]
image_tokens_sequence: list[LongTensor] = []
image_tokens_sequence: List[LongTensor] = []
for i in range(self.sample_token_count):
probs, attention_state = self.decode_step(
attention_mask = attention_mask,

3
min_dalle/models/dalle_bart_encoder.py
View File

@@ -1,3 +1,4 @@
from typing import List
import torch
from torch import nn, BoolTensor, FloatTensor, LongTensor
torch.set_grad_enabled(False)
@@ -120,7 +121,7 @@ class DalleBartEncoder(nn.Module):
super().__init__()
self.embed_tokens = nn.Embedding(text_vocab_count, embed_count)
self.embed_positions = nn.Embedding(text_token_count, embed_count)
self.layers: list[EncoderLayer] = nn.ModuleList([
self.layers: List[EncoderLayer] = nn.ModuleList([
EncoderLayer(
embed_count = embed_count,
head_count = attention_head_count,