Update README.md

README.md

@@ -46,32 +46,41 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("ruslanmv/Medical-Llama3-8B")
 model = AutoModelForCausalLM.from_pretrained("ruslanmv/Medical-Llama3-8B").to("cuda")  # If using GPU
-
-# Function to format and generate response with prompt engineering
+# Function to format and generate response with prompt engineering using a chat template
 def askme(question):
-    medical_prompt = """You are an AI Medical Assistant trained on a vast dataset of health information. Below is a medical question:
-
-    Question: {}
-
-    Please provide an informative and comprehensive answer:
-
-    Answer: """.format(question)
-
-    EOS_TOKEN = tokenizer.eos_token
-
-    def format_prompt(question):
-        return medical_prompt + question + EOS_TOKEN
-
-    inputs = tokenizer(format_prompt(question), return_tensors="pt").to("cuda")  # If using GPU
-    outputs = model.generate(**inputs, max_new_tokens=64, use_cache=True)  # Adjust max_new_tokens for longer responses
+    sys_message = ''' 
+    You are an AI Medical Assistant trained on a vast dataset of health information. Please be thorough and
+    provide an informative answer. If you don't know the answer to a specific medical inquiry, advise seeking professional help.
+    '''
+    
+    # Create messages structured for the chat template
+    messages = [{"role": "system", "content": sys_message}, {"role": "user", "content": question}]
+
+    # Applying chat template
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    outputs = model.generate(**inputs, max_new_tokens=100, use_cache=True)  # Adjust max_new_tokens for longer responses
+
+    # Extract and return the generated text
     answer = tokenizer.batch_decode(outputs)[0].strip()
     return answer
 
 # Example usage
-question = "What should I do to reduce my weight gained due to genetic hypothyroidism?"
+# - Context: First describe your problem.
+# - Question: Then make the question.
+question = '''
+I'm a 35-year-old male and for the past few months, I've been experiencing fatigue, increased sensitivity to cold, and dry, itchy skin.  
+
+Could these symptoms be related to hypothyroidism? 
+If so, what steps should I take to get a proper diagnosis and discuss treatment options?
+'''
 print(askme(question))
 ```
+the type of answer is :
+```
+Hello, I can understand your concern. I would suggest you to get a blood test done for TSH, T3, T4, and TPO antibodies. If the TSH is high, then you can be diagnosed with hypothyroidism. If the TSH is normal, then you can be diagnosed with Hashimoto's thyroiditis. In both cases, you will need to take levothyroxine. Hope I have answered your query. Let me know if I can assist you
 
+```
 **Important Note**
 
 This model is intended for informational purposes only and should not be used as a substitute for professional medical advice. Always consult with a qualified healthcare provider for any medical concerns.