Rearrange code

app.py

@@ -8,11 +8,6 @@ import spaces
 
 @spaces.GPU
 def main():
-    if torch.cuda.is_available():
-        device = torch.device('cuda')
-    else:
-        device = torch.device('cpu')
-
     n_layers = 12
     n_heads = 12
     embed_dim = 768
@@ -23,12 +18,15 @@ def main():
 
     tokenizer = tiktoken.encoding_for_model('gpt2')
     model_path = hf_hub_download('cbspace/gpt', 'model.safetensors')
-    state_dict = load_file(model_path)
-    model = GPTModel(device, n_layers, n_heads, embed_dim, ffn_dim, n_vocab, max_seq_len, dropout)
-    model.load_state_dict(state_dict, strict=False)
-    model.eval()
 
     def generate(prompt,out_tokens):
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        state_dict = load_file(model_path)
+
+        model = GPTModel(device, n_layers, n_heads, embed_dim, ffn_dim, n_vocab, max_seq_len, dropout)
+        model.load_state_dict(state_dict, strict=False)
+        model.eval()
+    
         inputs = tokenizer.encode(prompt)
         outputs = model.generate(inputs, int(out_tokens), greedy=False)
         return tokenizer.decode(outputs) 
@@ -37,9 +35,9 @@ def main():
         generate,
         [
             gr.Textbox(label='Prompt', lines=3), 
-            gr.Number(label='Output Tokens', value=32)
+            gr.Number(label='Output Tokens', value=100)
         ],
-        gr.Textbox(label='Output', lines=3)
+        gr.Textbox(label='Output', lines=6)
     )
     app.launch()