Update app.py

app.py

@@ -43,8 +43,11 @@ def chat(history, system_message, max_tokens, temperature, top_p, top_k, repetit
     # Apply the chat template
     gen_input = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True)
     
-    # Extract input_ids
-    input_ids = gen_input['input_ids']
+    # Debug: Print the shape of gen_input
+    print("Shape of gen_input:", gen_input.shape if hasattr(gen_input, 'shape') else type(gen_input))
+    
+    # Extract input_ids based on the type of gen_input
+    input_ids = gen_input['input_ids'] if isinstance(gen_input, dict) else gen_input
     
     # Generate the output
     output = model.generate(input_ids=input_ids, temperature=temperature, do_sample=True, top_p=top_p, top_k=top_k, max_length=max_tokens)
@@ -58,6 +61,7 @@ def chat(history, system_message, max_tokens, temperature, top_p, top_k, repetit
     return history, history, ""
 
 
+
 start_message = ""
 
 CSS ="""