Upload 2 files

app.py

@@ -1,13 +1,15 @@
 import streamlit as st
 import time
+import torch
 
 from better_transformer import *
+torch.set_default_device(torch.device("cuda"))
 
 def main():
 
     # Enable CUDA if available and load in tokenizer
-    device = torch.device('cuda') #if torch.cuda.is_available() else 'cpu')
-    tokenizer, EMPTY_TOKENS = load_tokenizer('cuda')
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    tokenizer, EMPTY_TOKENS = load_tokenizer(device)
 
     st.title("Scaling Transformers")
     st.subheader("UCLA DSU Project, Fall 2023")
@@ -37,7 +39,7 @@ def main():
             """
         )
     ## Default values for advanced settings
-    user_seed = 27 # Remove (set to None) if we're not rigging the "random" demo
+    user_seed = 27 # Remove if we're not rigging the "random" demo
     generation_method = "top-k"
     specified_k = 5
     specified_nucleus = 0.5
@@ -59,14 +61,17 @@ def main():
 
         max_tokens = st.slider('Max Tokens Generated:', 100, 500, 400)
 
-
+    ## Settings Clean up
+    if not user_seed:
+        user_seed = 7
 
 
 
 
     # 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, 'cuda')
+    model = load_big_model(tokenizer, device)
+    model.to('cuda')
 
 
 
@@ -83,7 +88,7 @@ def main():
         
 
         with st.spinner(""):
-            result = generate(model, tokenizer, 'cuda', method=generation_method, k=specified_k, 
+            result = generate(model, tokenizer, device, 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

@@ -9,8 +9,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers import AutoTokenizer
+from torchinfo import summary
 
 MODEL_FILE = r'bt_8_LAYERs_100_DATA_PCT_768_EMBD_DIM_epoch_10.pt' ##place model file in same directory as app.py
+torch.set_default_device(torch.device("cuda"))
 
 # Better Transformer Class –––––––––––––––––––––––––––––––––––––––––––––––
 
@@ -139,7 +141,7 @@ class BetterTransformer(nn.Module):
         self.seq_length = seq_length
         self.pad_idx = pad_idx
         self.eos_token_id = eos_token_id
-        self.device = 'cuda'
+        self.device = device
         self.init_params()
 
     # optional weight initialization (e.g. Xavier uniform)
@@ -294,7 +296,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('cuda')
+    EMPTY_TOKENS = torch.full((1,1), tokenizer.bos_token_id, dtype=torch.long).to(device)
     return tokenizer, EMPTY_TOKENS
 
 
@@ -307,10 +309,10 @@ def load_big_model(tokenizer, device):
     VOCAB_SIZE = 50258
     SEQ_LENGTH = 384
 
-    model = BetterTransformer(VOCAB_SIZE, SEQ_LENGTH, N_EMBD, N_HEAD, N_LAYER, tokenizer.pad_token_id, tokenizer.eos_token_id, device='cuda')
+    model = BetterTransformer(VOCAB_SIZE, SEQ_LENGTH, N_EMBD, N_HEAD, N_LAYER, tokenizer.pad_token_id, tokenizer.eos_token_id, device=device)
     model.init_params()
     path = MODEL_FILE
-    model.load_state_dict(torch.load(path, map_location='cuda')["model_state_dict"])
+    model.load_state_dict(torch.load(path, map_location=device)["model_state_dict"])
 
     return model
     
@@ -360,9 +362,9 @@ def generate(model, tokenizer, device, method=None, k=None,
 
     if cond != "":
 
-        cond_tokens = tokenizer(cond).input_ids
+        cond_tokens = tokenizer(cond).input_ids.to(device) ## HERE???
 
-        gen_tokens = model.generate(torch.tensor(cond_tokens).unsqueeze(0).long().to('cuda'),
+        gen_tokens = model.generate(torch.tensor(cond_tokens).unsqueeze(0).long().to(device),
                                     method=method, k=k, p_nucleus=p_nucleus, temp=temp,
                                     max_new_tokens=max_new_tokens)[0]
 
@@ -378,7 +380,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('cuda')
+        empty_tokens = torch.full((1,1), tokenizer.bos_token_id, dtype=torch.long).to(device)
 
         res = tokenizer.batch_decode(model.generate(empty_tokens,
                                                     method=method, k=k,