Update app.py

app.py

@@ -54,129 +54,129 @@ class CFG():
     num_workers=1
 
 if st.button('predict'):
-    st.progress(0)
-    
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    
-    
-    def seed_everything(seed=42):
-        random.seed(seed)
-        os.environ['PYTHONHASHSEED'] = str(seed)
-        np.random.seed(seed)
-        torch.manual_seed(seed)
-        torch.cuda.manual_seed(seed)
-        torch.backends.cudnn.deterministic = True
-    seed_everything(seed=CFG.seed)      
-    
-    CFG.tokenizer = AutoTokenizer.from_pretrained(CFG.model_name_or_path, return_tensors='pt')
-    
-    def prepare_input(cfg, text):
-        inputs = cfg.tokenizer(text, add_special_tokens=True, max_length=CFG.max_len, padding='max_length', return_offsets_mapping=False, truncation=True, return_attention_mask=True)
-        for k, v in inputs.items():
-            inputs[k] = torch.tensor(v, dtype=torch.long)
-        
-        return inputs
-    
-    class TestDataset(Dataset):
-        def __init__(self, cfg, df):
-            self.cfg = cfg
-            self.inputs = df['input'].values
-            
-        def __len__(self):
-            return len(self.inputs)
+    with st.spinner('Now processing. This process takes about 30 seconds per 10 reactions.'):
+        
+        
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        
+        def seed_everything(seed=42):
+            random.seed(seed)
+            os.environ['PYTHONHASHSEED'] = str(seed)
+            np.random.seed(seed)
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed(seed)
+            torch.backends.cudnn.deterministic = True
+        seed_everything(seed=CFG.seed)      
+        
+        CFG.tokenizer = AutoTokenizer.from_pretrained(CFG.model_name_or_path, return_tensors='pt')
         
-        def __getitem__(self, item):
-            inputs = prepare_input(self.cfg, self.inputs[item])
+        def prepare_input(cfg, text):
+            inputs = cfg.tokenizer(text, add_special_tokens=True, max_length=CFG.max_len, padding='max_length', return_offsets_mapping=False, truncation=True, return_attention_mask=True)
+            for k, v in inputs.items():
+                inputs[k] = torch.tensor(v, dtype=torch.long)
             
             return inputs
         
-    
-    class RegressionModel(nn.Module):
-        def __init__(self, cfg, config_path=None, pretrained=False):
-            super().__init__()
-            self.cfg = cfg
-            if config_path is None:
-                self.config = AutoConfig.from_pretrained(cfg.pretrained_model_name_or_path, output_hidden_states=True)
-            else:
-                self.config = torch.load(config_path)
-            if pretrained:
-                if 't5' in cfg.model:
-                    self.model = T5EncoderModel.from_pretrained(CFG.pretrained_model_name_or_path)
+        class TestDataset(Dataset):
+            def __init__(self, cfg, df):
+                self.cfg = cfg
+                self.inputs = df['input'].values
+                
+            def __len__(self):
+                return len(self.inputs)
+            
+            def __getitem__(self, item):
+                inputs = prepare_input(self.cfg, self.inputs[item])
+                
+                return inputs
+            
+        
+        class RegressionModel(nn.Module):
+            def __init__(self, cfg, config_path=None, pretrained=False):
+                super().__init__()
+                self.cfg = cfg
+                if config_path is None:
+                    self.config = AutoConfig.from_pretrained(cfg.pretrained_model_name_or_path, output_hidden_states=True)
                 else:
-                    self.model = AutoModel.from_pretrained(CFG.pretrained_model_name_or_path)
-            else:
-                if 't5' in cfg.model:
-                    self.model = T5EncoderModel.from_pretrained('sagawa/ZINC-t5')
+                    self.config = torch.load(config_path)
+                if pretrained:
+                    if 't5' in cfg.model:
+                        self.model = T5EncoderModel.from_pretrained(CFG.pretrained_model_name_or_path)
+                    else:
+                        self.model = AutoModel.from_pretrained(CFG.pretrained_model_name_or_path)
                 else:
-                    self.model = AutoModel.from_config(self.config)
-            self.model.resize_token_embeddings(len(cfg.tokenizer))
-            self.fc_dropout1 = nn.Dropout(cfg.fc_dropout)
-            self.fc1 = nn.Linear(self.config.hidden_size, self.config.hidden_size)
-            self.fc_dropout2 = nn.Dropout(cfg.fc_dropout)
-            self.fc2 = nn.Linear(self.config.hidden_size, 1)
+                    if 't5' in cfg.model:
+                        self.model = T5EncoderModel.from_pretrained('sagawa/ZINC-t5')
+                    else:
+                        self.model = AutoModel.from_config(self.config)
+                self.model.resize_token_embeddings(len(cfg.tokenizer))
+                self.fc_dropout1 = nn.Dropout(cfg.fc_dropout)
+                self.fc1 = nn.Linear(self.config.hidden_size, self.config.hidden_size)
+                self.fc_dropout2 = nn.Dropout(cfg.fc_dropout)
+                self.fc2 = nn.Linear(self.config.hidden_size, 1)
+                
+            def forward(self, inputs):
+                outputs = self.model(**inputs)
+                last_hidden_states = outputs[0]
+                output = self.fc1(self.fc_dropout1(last_hidden_states)[:, 0, :].view(-1, self.config.hidden_size))
+                output = self.fc2(self.fc_dropout2(output))
+                return output
+            
+        
+            
+        def inference_fn(test_loader, model, device):
+            preds = []
+            model.eval()
+            model.to(device)
+            tk0 = enumerate(test_loader)
+            for i, inputs in tk0:
+                for k, v in inputs.items():
+                    inputs[k] = v.to(device)
+                with torch.no_grad():
+                    y_preds = model(inputs)
+                preds.append(y_preds.to('cpu').numpy())
+            predictions = np.concatenate(preds)
+            return predictions
+        
+        model = RegressionModel(CFG, config_path=CFG.model_name_or_path + '/config.pth', pretrained=False)
+        state = torch.load(CFG.model_name_or_path + '/ZINC-t5_best.pth', map_location=torch.device('cpu'))
+        model.load_state_dict(state)
+        
+        
+        if CFG.uploaded_file is not None:
+            test_ds = pd.read_csv(CFG.uploaded_file)
+                
+            test_dataset = TestDataset(CFG, test_ds)
+            test_loader = DataLoader(test_dataset,
+                                     batch_size=CFG.batch_size,
+                                     shuffle=False,
+                                     num_workers=CFG.num_workers, pin_memory=True, drop_last=False)
+        
+        
+            prediction = inference_fn(test_loader, model, device)
+        
+            test_ds['prediction'] = prediction*100
+            test_ds['prediction'] = test_ds['prediction'].clip(0, 100)
+            csv = test_ds.to_csv(index=False)
+            st.download_button(
+                label="Download data as CSV",
+                data=csv,
+                file_name='output.csv',
+                mime='text/csv'
+            )
+        
+        else:
+            CFG.batch_size=1
+            test_ds = pd.DataFrame.from_dict({'input': CFG.data}, orient='index').T
+            test_dataset = TestDataset(CFG, test_ds)
+            test_loader = DataLoader(test_dataset,
+                                     batch_size=CFG.batch_size,
+                                     shuffle=False,
+                                     num_workers=CFG.num_workers, pin_memory=True, drop_last=False)
             
-        def forward(self, inputs):
-            outputs = self.model(**inputs)
-            last_hidden_states = outputs[0]
-            output = self.fc1(self.fc_dropout1(last_hidden_states)[:, 0, :].view(-1, self.config.hidden_size))
-            output = self.fc2(self.fc_dropout2(output))
-            return output
-        
-    
-        
-    def inference_fn(test_loader, model, device):
-        preds = []
-        model.eval()
-        model.to(device)
-        tk0 = enumerate(test_loader)
-        for i, inputs in tk0:
-            for k, v in inputs.items():
-                inputs[k] = v.to(device)
-            with torch.no_grad():
-                y_preds = model(inputs)
-            st.progress((i+1)*CFG.batch_size)
-            preds.append(y_preds.to('cpu').numpy())
-        predictions = np.concatenate(preds)
-        return predictions
-    
-    model = RegressionModel(CFG, config_path=CFG.model_name_or_path + '/config.pth', pretrained=False)
-    state = torch.load(CFG.model_name_or_path + '/ZINC-t5_best.pth', map_location=torch.device('cpu'))
-    model.load_state_dict(state)
-    
-    
-    if CFG.uploaded_file is not None:
-        test_ds = pd.read_csv(CFG.uploaded_file)
             
-        test_dataset = TestDataset(CFG, test_ds)
-        test_loader = DataLoader(test_dataset,
-                                 batch_size=CFG.batch_size,
-                                 shuffle=False,
-                                 num_workers=CFG.num_workers, pin_memory=True, drop_last=False)
-    
-    
-        prediction = inference_fn(test_loader, model, device)
-    
-        test_ds['prediction'] = prediction*100
-        test_ds['prediction'] = test_ds['prediction'].clip(0, 100)
-        csv = test_ds.to_csv(index=False)
-        st.download_button(
-            label="Download data as CSV",
-            data=csv,
-            file_name='output.csv',
-            mime='text/csv'
-        )
-    
-    else:
-        CFG.batch_size=1
-        test_ds = pd.DataFrame.from_dict({'input': CFG.data}, orient='index').T
-        test_dataset = TestDataset(CFG, test_ds)
-        test_loader = DataLoader(test_dataset,
-                                 batch_size=CFG.batch_size,
-                                 shuffle=False,
-                                 num_workers=CFG.num_workers, pin_memory=True, drop_last=False)
-        
-        
-        prediction = inference_fn(test_loader, model, device)
-        prediction = max(min(prediction[0][0]*100, 100), 0)
-        st.text('yiled: '+ str(prediction))
- 
+            prediction = inference_fn(test_loader, model, device)
+            prediction = max(min(prediction[0][0]*100, 100), 0)
+            st.text('yiled: '+ str(prediction))
+