Spaces:
Sleeping
Sleeping
| from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline | |
| from typing import Dict, List, Tuple | |
| import uuid | |
| MODEL_NAME = "BioMistral/BioMistral-7B" | |
| print("🔹 Loading BioMistral model... This may take a while on first run.") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_NAME, | |
| device_map="auto", | |
| torch_dtype="auto" | |
| ) | |
| chat_pipeline = pipeline( | |
| "text-generation", | |
| model=model, | |
| tokenizer=tokenizer, | |
| max_new_tokens=512, | |
| temperature=0.7, | |
| top_p=0.9 | |
| ) | |
| # Store conversation history per session | |
| conversation_sessions: Dict[str, List[Dict[str, str]]] = {} | |
| SYSTEM_PROMPT = ( | |
| "You are MediScope AI, a medical assistant that helps patients understand " | |
| "their symptoms in simple, safe, and educational language. " | |
| "Always encourage professional consultation for serious conditions." | |
| ) | |
| def get_or_create_session(session_id: str) -> List[Dict[str, str]]: | |
| """Get existing session or create a new one.""" | |
| if session_id not in conversation_sessions: | |
| conversation_sessions[session_id] = [] | |
| return conversation_sessions[session_id] | |
| def build_conversation_prompt(history: List[Dict[str, str]], user_prompt: str) -> str: | |
| """Build the full conversation prompt from history and new user message.""" | |
| prompt_parts = [SYSTEM_PROMPT] | |
| # Add conversation history | |
| for msg in history: | |
| if msg["role"] == "user": | |
| prompt_parts.append(f"User: {msg['content']}") | |
| elif msg["role"] == "assistant": | |
| prompt_parts.append(f"Assistant: {msg['content']}") | |
| # Add current user message | |
| prompt_parts.append(f"User: {user_prompt}") | |
| prompt_parts.append("Assistant:") | |
| return "\n\n".join(prompt_parts) | |
| def chat_with_biomistral(user_prompt: str, session_id: str = None) -> Tuple[str, str]: | |
| """ | |
| Chat with BioMistral model, maintaining conversation history. | |
| Args: | |
| user_prompt: The user's message | |
| session_id: Optional session ID. If None, a new session is created. | |
| Returns: | |
| tuple: (response_text, session_id) | |
| """ | |
| # Generate or use provided session ID | |
| if session_id is None: | |
| session_id = str(uuid.uuid4()) | |
| # Get conversation history for this session | |
| history = get_or_create_session(session_id) | |
| # Build the full conversation prompt | |
| full_prompt = build_conversation_prompt(history, user_prompt) | |
| # Generate response | |
| response = chat_pipeline(full_prompt)[0]["generated_text"] | |
| # Extract only the assistant's reply (everything after the last "Assistant:") | |
| reply = response.split("Assistant:")[-1].strip() | |
| # Update conversation history | |
| history.append({"role": "user", "content": user_prompt}) | |
| history.append({"role": "assistant", "content": reply}) | |
| return reply, session_id | |
| def clear_session(session_id: str) -> bool: | |
| """Clear conversation history for a session.""" | |
| if session_id in conversation_sessions: | |
| conversation_sessions[session_id] = [] | |
| return True | |
| return False | |