Spaces:
Running on Zero
Running on Zero
| import os | |
| import tempfile | |
| import shutil | |
| import numpy as np | |
| from tqdm.notebook import tqdm | |
| import requests | |
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| from torch.utils.data import DataLoader, TensorDataset | |
| from torch.optim.lr_scheduler import CosineAnnealingLR | |
| from sklearn.feature_extraction.text import HashingVectorizer | |
| import logging | |
| class ResidualBlock(nn.Module): | |
| def __init__(self, hidden_size, dropout_prob): | |
| super(ResidualBlock, self).__init__() | |
| self.block = nn.Sequential( | |
| nn.Linear(hidden_size, hidden_size), | |
| nn.LayerNorm(hidden_size), | |
| nn.ReLU(), | |
| nn.Dropout(dropout_prob), | |
| nn.Linear(hidden_size, hidden_size), | |
| nn.LayerNorm(hidden_size), | |
| ) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| residual = x | |
| out = self.block(x) | |
| out += residual # Skip connection | |
| return self.relu(out) | |
| class DeepNeuralNetwork(nn.Module): | |
| def __init__(self, input_size, num_layers=10, hidden_size=4096, dropout_prob=0.2): | |
| super(DeepNeuralNetwork, self).__init__() | |
| # First layer | |
| self.input_layer = nn.Sequential( | |
| nn.Linear(input_size, hidden_size), | |
| nn.LayerNorm(hidden_size), | |
| nn.ReLU(), | |
| nn.Dropout(dropout_prob), | |
| ) | |
| # Residual blocks | |
| self.residual_blocks = nn.ModuleList() | |
| for i in range(num_layers - 2): | |
| self.residual_blocks.append(ResidualBlock(hidden_size, dropout_prob)) | |
| # Output layer | |
| self.output_layer = nn.Linear(hidden_size, 1) | |
| def forward(self, x): | |
| x = self.input_layer(x) | |
| for block in self.residual_blocks: | |
| x = block(x) | |
| return self.output_layer(x) | |
| Y_STD = 1.0328539609909058 | |
| Y_MEAN = 4.434937953948975 | |
| class DeepNeuralNetworkInference: | |
| def __init__(self): | |
| self.vectorizer = None | |
| self.model = None | |
| self.device = None | |
| np.random.seed(42) | |
| torch.manual_seed(42) | |
| torch.cuda.manual_seed(42) | |
| def setup(self): | |
| self.vectorizer = HashingVectorizer(n_features=5000, stop_words="english", binary=True) | |
| self.model = DeepNeuralNetwork(5000) | |
| if torch.cuda.is_available(): | |
| self.device = torch.device("cuda") | |
| elif torch.backends.mps.is_available(): | |
| self.device = torch.device("mps") | |
| else: | |
| self.device = torch.device("cpu") | |
| logging.info(f"Neural Network is using {self.device}") | |
| self.model.to(self.device) | |
| def load(self, path): | |
| self.model.load_state_dict(torch.load(path, map_location=self.device)) | |
| self.model.to(self.device) | |
| def inference(self, text): | |
| self.model.eval() | |
| with torch.no_grad(): | |
| vector = self.vectorizer.transform([text]) | |
| vector = torch.FloatTensor(vector.toarray()).to(self.device) | |
| pred = self.model(vector)[0] | |
| result = torch.exp(pred * Y_STD + Y_MEAN) - 1 | |
| result = result.item() | |
| return max(0, result) | |