Update app.py

app.py

@@ -53,15 +53,21 @@ def _launch_demo(args, model, tokenizer, config):
     def predict(_query, _chatbot, _task_history):
         print(f"User: {_parse_text(_query)}")
         _chatbot.append((_parse_text(_query), ""))
-        input_ids = input_ids.to('cuda')
-        attention_mask = torch.ones(input_ids.shape).to('cuda')
-        pad_token_id = tokenizer.eos_token_id
+        
         # Tokenize the input
         input_ids = tokenizer.encode(_query, return_tensors='pt')
         print("Input IDs:", input_ids)
+        
+        # Move input_ids to CUDA if available
+        input_ids = input_ids.to('cuda')
+        
+        # Generate attention_mask
+        attention_mask = torch.ones(input_ids.shape).to('cuda')
+        
         # Generate a response using the model
         generated_ids = model.generate(input_ids, max_length=300)
         print("Generated IDs:", generated_ids)
+        
         # Decode the generated IDs to text
         full_response = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
         
@@ -73,6 +79,7 @@ def _launch_demo(args, model, tokenizer, config):
         _task_history.append((_query, full_response))
         print(f"OpenHermes: {_parse_text(full_response)}")
 
+
     def regenerate(_chatbot, _task_history):
         if not _task_history:
             yield _chatbot