Update app.py

app.py

@@ -18,7 +18,6 @@ model_name = "ruslanmv/Medical-Llama3-8B"
 model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
-
 @spaces.GPU
 def askme(symptoms, question):
     sys_message = '''\
@@ -28,14 +27,47 @@ def askme(symptoms, question):
     content = symptoms + " " + question
     messages = [{"role": "system", "content": sys_message}, {"role": "user", "content": content}]
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)  # Ensure inputs are on CUDA device
-    outputs = model.generate(**inputs, max_new_tokens=200, use_cache=True)
-    response_text = tokenizer.batch_decode(outputs)[0].strip()
+   # inputs = tokenizer(prompt, return_tensors="pt").to(device)  # Ensure inputs are on CUDA device
+   # outputs = model.generate(**inputs, max_new_tokens=200, use_cache=True)
+   # response_text = tokenizer.batch_decode(outputs)[0].strip()
+
+
+
+    # Tokenize all prompts and batch them
+    tokenized_inputs = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True).to(device)
+    outputs = model.generate(**tokenized_inputs, max_new_tokens=200, use_cache=True)
+
+    # Decode responses
+    response_texts = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+    
+    
     # Remove system messages and content
     #response_text = response_text.replace(sys_message, "").replace(content, "").strip()
     # Extract only the assistant's response
-    return response_text
+    #assistant_response = response_text.split("<|im_start|>assistant")[1].strip().replace('<|im_end', '')
+    # Extract only the assistant's response
+    
+    
+    
+    
+       # Extract assistant's responses
+    assistant_responses = []
+    for response_text in response_texts:
+        assistant_response = response_text.split("assistant")[1].strip().replace('<|im_end', '')
+        assistant_responses.append(assistant_response)
 
+    return assistant_responses
+    
+    
+    
+    
+    
+    
+    
+    
+    
+    return assistant_response
 # Example usage
 symptoms = '''\
 I'm a 35-year-old male and for the past few months, I've been experiencing fatigue,