Upload 2 files

app.py

@@ -66,7 +66,7 @@ def main():
 
     # model_version = st.radio("Which model would you like to use?", ["smoll", "beeg"])
     # small_model = load_casey_model(tokenizer, device)
-    model = load_big_model(tokenizer, device)
+    model = load_big_model(tokenizer, 'cuda')
 
 
 
@@ -83,7 +83,7 @@ def main():
         
 
         with st.spinner(""):
-            result = generate(model, tokenizer, device, method=generation_method, k=specified_k, 
+            result = generate(model, tokenizer, 'cuda', method=generation_method, k=specified_k, 
                             p_nucleus=specified_nucleus, temp=specified_temperature, max_new_tokens=max_tokens, 
                             cond=user_input, deterministic=user_seed)
 

better_transformer.py

@@ -139,9 +139,8 @@ class BetterTransformer(nn.Module):
         self.seq_length = seq_length
         self.pad_idx = pad_idx
         self.eos_token_id = eos_token_id
-        self.device = device
+        self.device = 'cuda'
         self.init_params()
-        st.title(f"Device: {device}")
 
     # optional weight initialization (e.g. Xavier uniform)
     def init_params(self, default_initialization=False):
@@ -295,7 +294,7 @@ def load_tokenizer(device):
     tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
     if tokenizer.pad_token is None:
         tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-    EMPTY_TOKENS = torch.full((1,1), tokenizer.bos_token_id, dtype=torch.long).to(device)
+    EMPTY_TOKENS = torch.full((1,1), tokenizer.bos_token_id, dtype=torch.long).to('cuda')
     return tokenizer, EMPTY_TOKENS
 
 
@@ -363,7 +362,7 @@ def generate(model, tokenizer, device, method=None, k=None,
 
         cond_tokens = tokenizer(cond).input_ids
 
-        gen_tokens = model.generate(torch.tensor(cond_tokens).unsqueeze(0).long().to(device),
+        gen_tokens = model.generate(torch.tensor(cond_tokens).unsqueeze(0).long().to('cuda'),
                                     method=method, k=k, p_nucleus=p_nucleus, temp=temp,
                                     max_new_tokens=max_new_tokens)[0]
 
@@ -379,7 +378,7 @@ def generate(model, tokenizer, device, method=None, k=None,
         
 
     else:
-        empty_tokens = torch.full((1,1), tokenizer.bos_token_id, dtype=torch.long).to(device)
+        empty_tokens = torch.full((1,1), tokenizer.bos_token_id, dtype=torch.long).to('cuda')
 
         res = tokenizer.batch_decode(model.generate(empty_tokens,
                                                     method=method, k=k,