Update app.py

app.py

@@ -20,7 +20,11 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
-# Set the chat template
+# Set pad_token_id to eos_token_id if None
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+
+# Define the chat template
 chat_template = """<|im_start|>system
 {system}
 <|im_end|>
@@ -31,7 +35,7 @@ chat_template = """<|im_start|>system
 """
 tokenizer.chat_template = chat_template
 
-# Define the askme function
+# Function to generate a response
 def askme(question):
     sys_message = """ 
     You are an AI Medical Assistant trained on a vast dataset of health information. Please be thorough and
@@ -47,10 +51,11 @@ def askme(question):
     # Generate response
     outputs = model.generate(**inputs, max_new_tokens=100, use_cache=True)
     
-    # Extract and return the generated text
+    # Decode and clean up the response
     response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    answer = response_text.split("<|im_start|>assistant")[-1].strip()
-    return answer
+    if "<|im_start|>assistant" in response_text:
+        response_text = response_text.split("<|im_start|>assistant")[-1].strip()
+    return response_text
 
 # Example usage
 question = """