Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,795 Bytes
6e15dcc 031a3f5 6e15dcc 031a3f5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
import spaces
from .model import ModelManager
from .memory import MedicalMemoryManager
from .prompts import CONSULTATION_PROMPT, MEDICINE_PROMPT
model_manager = ModelManager()
memory_manager = MedicalMemoryManager()
conversation_turns = 0
def build_me_llama_prompt(system_prompt, history, user_input):
memory_context = memory_manager.get_memory_context()
enhanced_system_prompt = f"{system_prompt}\n\nPrevious conversation context:\n{memory_context}"
prompt = f"<s>[INST] <<SYS>>\n{enhanced_system_prompt}\n<</SYS>>\n\n"
recent_history = history[-3:] if len(history) > 3 else history
for user_msg, assistant_msg in recent_history:
prompt += f"{user_msg} [/INST] {assistant_msg} </s><s>[INST] "
prompt += f"{user_input} [/INST] "
return prompt
@spaces.GPU
def respond(message, chat_history):
global conversation_turns
conversation_turns += 1
if conversation_turns < 4:
prompt = build_me_llama_prompt(CONSULTATION_PROMPT, chat_history, message)
response = model_manager.generate(prompt)
memory_manager.add_interaction(message, response)
chat_history.append((message, response))
return "", chat_history
else:
patient_summary = memory_manager.get_patient_summary()
memory_context = memory_manager.get_memory_context()
summary_prompt = build_me_llama_prompt(
CONSULTATION_PROMPT + "\n\nNow provide a comprehensive summary based on all the information gathered. Include when professional care may be needed.",
chat_history,
message
)
summary = model_manager.generate(summary_prompt)
full_patient_info = f"Patient Summary: {patient_summary}\n\nDetailed Summary: {summary}"
med_prompt = f"<s>[INST] {MEDICINE_PROMPT.format(patient_info=full_patient_info, memory_context=memory_context)} [/INST] "
medicine_suggestions = model_manager.generate(med_prompt, max_new_tokens=300)
final_response = (
f"**COMPREHENSIVE MEDICAL SUMMARY:**\n{summary}\n\n"
f"**MEDICATION AND HOME CARE SUGGESTIONS:**\n{medicine_suggestions}\n\n"
f"**PATIENT CONTEXT SUMMARY:**\n{patient_summary}\n\n"
f"**DISCLAIMER:** This is AI-generated advice for informational purposes only. Please consult a licensed healthcare provider for proper medical diagnosis and treatment."
)
memory_manager.add_interaction(message, final_response)
chat_history.append((message, final_response))
return "", chat_history
def reset_chat():
global conversation_turns
conversation_turns = 0
memory_manager.reset_session()
reset_msg = "New consultation started. Please tell me about your symptoms or health concerns."
return [(None, reset_msg)], "" |