Refactor code

app.py

@@ -6,13 +6,13 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 model = torch.load('saved_model.pth', map_location= torch.device(device), weights_only=False)
 
-def generate_text(context, num_of_tokens, temperature=1.0):
-    if context == None or context == '':
-        idx = torch.zeros((1,1), dtype=torch.long)
-    else:
-        idx = torch.tensor(encode(context), dtype=torch.long).unsqueeze(0)
+# def generate_text(context, num_of_tokens, temperature=1.0):
+#     if context == None or context == '':
+#         idx = torch.zeros((1,1), dtype=torch.long)
+#     else:
+#         idx = torch.tensor(encode(context), dtype=torch.long).unsqueeze(0)
     
-    yield model.generate(idx, max_new_tokens=num_of_tokens,temperature=temperature)[0].tolist() 
+#     yield model.generate(idx, max_new_tokens=num_of_tokens,temperature=temperature)[0].tolist() 
     
 
 
@@ -29,8 +29,12 @@ with gr.Blocks() as demo:
         context,
         num_of_tokens,tmp
     ]
+    if context == None or context == '':
+        idx = torch.zeros((1,1), dtype=torch.long)
+    else:
+        idx = torch.tensor(encode(context), dtype=torch.long).unsqueeze(0)
     generate_btn = gr.Button(value="Generate")
     outputs = [gr.Textbox(label= "Generated text: ")]
-    generate_btn.click(fn = generate_text, inputs= inputs, outputs= outputs)
+    generate_btn.click(fn = model.generate(idx, max_new_tokens=num_of_tokens,temperature=tmp)[0].tolist() , inputs= inputs, outputs= outputs)
     
 demo.launch()