wooyeolbaek commited on
Commit
c47b6ff
1 Parent(s): 7971f77

Update app.py

Browse files

remove sos and eos tokens

Files changed (1) hide show
  1. app.py +4 -0
app.py CHANGED
@@ -28,6 +28,10 @@ def inference(prompt):
28
  net_attn_maps = get_net_attn_map(image.size)
29
  net_attn_maps = resize_net_attn_map(net_attn_maps, image.size)
30
  net_attn_maps = return_net_attn_map(net_attn_maps, pipe.tokenizer, prompt)
 
 
 
 
31
 
32
  return image, net_attn_maps
33
 
 
28
  net_attn_maps = get_net_attn_map(image.size)
29
  net_attn_maps = resize_net_attn_map(net_attn_maps, image.size)
30
  net_attn_maps = return_net_attn_map(net_attn_maps, pipe.tokenizer, prompt)
31
+
32
+ # remove sos and eos
33
+ net_attn_maps = [attn_map for attn_map in net_attn_maps if attn_map[1].split('_')[-1] != "<<|startoftext|>>"]
34
+ net_attn_maps = [attn_map for attn_map in net_attn_maps if attn_map[1].split('_')[-1] != "<<|endoftext|>>"]
35
 
36
  return image, net_attn_maps
37