debisoft commited on
Commit
814c118
1 Parent(s): ee92395

1st commit!! :D

Browse files
Files changed (2) hide show
  1. app.py +136 -0
  2. requirements.txt +2 -0
app.py ADDED
@@ -0,0 +1,136 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, Tuple
2
+ from tqdm import tqdm
3
+ import torch
4
+ import torch.nn as nn
5
+ import torch.nn.functional as F
6
+ from torch.utils.data import DataLoader
7
+ from torchvision import models, transforms
8
+ from torchvision.utils import save_image, make_grid
9
+ import matplotlib.pyplot as plt
10
+ from matplotlib.animation import FuncAnimation, PillowWriter
11
+ import numpy as np
12
+ from IPython.display import HTML
13
+ from diffusion_utilities import *
14
+
15
+ openai.api_key = os.getenv('OPENAI_API_KEY')
16
+
17
+ class ContextUnet(nn.Module):
18
+ def __init__(self, in_channels, n_feat=256, n_cfeat=10, height=28): # cfeat - context features
19
+ super(ContextUnet, self).__init__()
20
+
21
+ # number of input channels, number of intermediate feature maps and number of classes
22
+ self.in_channels = in_channels
23
+ self.n_feat = n_feat
24
+ self.n_cfeat = n_cfeat
25
+ self.h = height #assume h == w. must be divisible by 4, so 28,24,20,16...
26
+
27
+ # Initialize the initial convolutional layer
28
+ self.init_conv = ResidualConvBlock(in_channels, n_feat, is_res=True)
29
+
30
+ # Initialize the down-sampling path of the U-Net with two levels
31
+ self.down1 = UnetDown(n_feat, n_feat) # down1 #[10, 256, 8, 8]
32
+ self.down2 = UnetDown(n_feat, 2 * n_feat) # down2 #[10, 256, 4, 4]
33
+
34
+ # original: self.to_vec = nn.Sequential(nn.AvgPool2d(7), nn.GELU())
35
+ self.to_vec = nn.Sequential(nn.AvgPool2d((4)), nn.GELU())
36
+
37
+ # Embed the timestep and context labels with a one-layer fully connected neural network
38
+ self.timeembed1 = EmbedFC(1, 2*n_feat)
39
+ self.timeembed2 = EmbedFC(1, 1*n_feat)
40
+ self.contextembed1 = EmbedFC(n_cfeat, 2*n_feat)
41
+ self.contextembed2 = EmbedFC(n_cfeat, 1*n_feat)
42
+
43
+ # Initialize the up-sampling path of the U-Net with three levels
44
+ self.up0 = nn.Sequential(
45
+ nn.ConvTranspose2d(2 * n_feat, 2 * n_feat, self.h//4, self.h//4),
46
+ nn.GroupNorm(8, 2 * n_feat), # normalize
47
+ nn.ReLU(),
48
+ )
49
+ self.up1 = UnetUp(4 * n_feat, n_feat)
50
+ self.up2 = UnetUp(2 * n_feat, n_feat)
51
+
52
+ # Initialize the final convolutional layers to map to the same number of channels as the input image
53
+ self.out = nn.Sequential(
54
+ nn.Conv2d(2 * n_feat, n_feat, 3, 1, 1), # reduce number of feature maps #in_channels, out_channels, kernel_size, stride=1, padding=0
55
+ nn.GroupNorm(8, n_feat), # normalize
56
+ nn.ReLU(),
57
+ nn.Conv2d(n_feat, self.in_channels, 3, 1, 1), # map to same number of channels as input
58
+ )
59
+
60
+ def forward(self, x, t, c=None):
61
+ """
62
+ x : (batch, n_feat, h, w) : input image
63
+ t : (batch, n_cfeat) : time step
64
+ c : (batch, n_classes) : context label
65
+ """
66
+ # x is the input image, c is the context label, t is the timestep, context_mask says which samples to block the context on
67
+
68
+ # pass the input image through the initial convolutional layer
69
+ x = self.init_conv(x)
70
+ # pass the result through the down-sampling path
71
+ down1 = self.down1(x) #[10, 256, 8, 8]
72
+ down2 = self.down2(down1) #[10, 256, 4, 4]
73
+
74
+ # convert the feature maps to a vector and apply an activation
75
+ hiddenvec = self.to_vec(down2)
76
+
77
+ # mask out context if context_mask == 1
78
+ if c is None:
79
+ c = torch.zeros(x.shape[0], self.n_cfeat).to(x)
80
+
81
+ # embed context and timestep
82
+ cemb1 = self.contextembed1(c).view(-1, self.n_feat * 2, 1, 1) # (batch, 2*n_feat, 1,1)
83
+ temb1 = self.timeembed1(t).view(-1, self.n_feat * 2, 1, 1)
84
+ cemb2 = self.contextembed2(c).view(-1, self.n_feat, 1, 1)
85
+ temb2 = self.timeembed2(t).view(-1, self.n_feat, 1, 1)
86
+ #print(f"uunet forward: cemb1 {cemb1.shape}. temb1 {temb1.shape}, cemb2 {cemb2.shape}. temb2 {temb2.shape}")
87
+
88
+
89
+ up1 = self.up0(hiddenvec)
90
+ up2 = self.up1(cemb1*up1 + temb1, down2) # add and multiply embeddings
91
+ up3 = self.up2(cemb2*up2 + temb2, down1)
92
+ out = self.out(torch.cat((up3, x), 1))
93
+ return out
94
+
95
+ # hyperparameters
96
+
97
+ # diffusion hyperparameters
98
+ timesteps = 500
99
+ beta1 = 1e-4
100
+ beta2 = 0.02
101
+
102
+ # network hyperparameters
103
+ device = torch.device("cuda:0" if torch.cuda.is_available() else torch.device('cpu'))
104
+ n_feat = 64 # 64 hidden dimension feature
105
+ n_cfeat = 5 # context vector is of size 5
106
+ height = 16 # 16x16 image
107
+ save_dir = './weights/'
108
+
109
+ # training hyperparameters
110
+ batch_size = 100
111
+ n_epoch = 32
112
+ lrate=1e-3
113
+
114
+ # construct DDPM noise schedule
115
+ b_t = (beta2 - beta1) * torch.linspace(0, 1, timesteps + 1, device=device) + beta1
116
+ a_t = 1 - b_t
117
+ ab_t = torch.cumsum(a_t.log(), dim=0).exp()
118
+ ab_t[0] = 1
119
+
120
+ # construct model
121
+ nn_model = ContextUnet(in_channels=3, n_feat=n_feat, n_cfeat=n_cfeat, height=height).to(device)
122
+
123
+ def greet(input):
124
+ prompt = f"""
125
+ Recommend complementary shop combinations which match well with the shop(s) described in the following text, which is delimited by triple backticks. Rank by synergy: \
126
+ Text: ```{input}```
127
+ """
128
+ response = prompt
129
+ return response
130
+
131
+ #iface = gr.Interface(fn=greet, inputs="text", outputs="text")
132
+ #iface.launch()
133
+
134
+ #iface = gr.Interface(fn=greet, inputs=[gr.Textbox(label="Text to find entities", lines=2)], outputs=[gr.HighlightedText(label="Text with entities")], title="NER with dslim/bert-base-NER", description="Find entities using the `dslim/bert-base-NER` model under the hood!", allow_flagging="never", examples=["My name is Andrew and I live in California", "My name is Poli and work at HuggingFace"])
135
+ iface = gr.Interface(fn=greet, inputs=[gr.Textbox(label="Co-Retailing Business")], outputs="text")
136
+ iface.launch()
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ openai
2
+ python-dotenv