File size: 5,584 Bytes
43e7ae4 |
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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
"""
AI Services Module
This module handles AI-powered features including
chat responses and structure generation.
"""
from huggingface_hub import InferenceClient
from rdkit import Chem
from rdkit.Chem import Draw
def respond(
message,
history: list[dict[str, str]],
system_message,
max_tokens,
temperature,
top_p,
hf_token: str,
):
"""
Enhanced drug discovery chatbot with molecular property integration.
"""
client = InferenceClient(token=hf_token, model="openai/gpt-oss-20b")
# Enhanced system message for drug discovery
enhanced_system_message = f"""{system_message}
You are an expert medicinal chemist and drug discovery specialist. You help researchers understand:
- Molecular properties and drug-likeness
- Structure-activity relationships (SAR)
- ADMET properties (Absorption, Distribution, Metabolism, Excretion, Toxicity)
- Drug design principles and optimization strategies
- Chemical synthesis and medicinal chemistry
When discussing molecules, consider their molecular weight, LogP, hydrogen bonding, and other key properties. Reference the molecular gallery below for visual context."""
messages = [{"role": "system", "content": enhanced_system_message}]
messages.extend(history)
messages.append({"role": "user", "content": message})
response = ""
for message in client.chat_completion(
messages,
max_tokens=max_tokens,
stream=True,
temperature=temperature,
top_p=top_p,
):
choices = message.choices
token = ""
if len(choices) and choices[0].delta.content:
token = choices[0].delta.content
response += token
yield response
def generate_ai_structures(message, history, selected_smiles, hf_token: str):
"""Generate new structures using AI based on user request and selected molecule."""
if not message.strip():
return history, []
client = InferenceClient(token=hf_token, model="openai/gpt-oss-20b")
# Enhanced system message for structure generation
system_message = f"""You are an expert medicinal chemist and drug discovery specialist. You help generate new chemical structures based on user requests and existing molecules.
Current selected molecule SMILES: {selected_smiles}
Your task is to:
1. Understand the user's request for structure modifications
2. Generate 3-6 new SMILES strings that meet their requirements
3. Provide brief explanations for each generated structure
4. Consider drug-likeness, ADMET properties, and medicinal chemistry principles
Format your response as:
STRUCTURE 1: [SMILES] - [Brief explanation]
STRUCTURE 2: [SMILES] - [Brief explanation]
etc.
Focus on practical, synthesizable structures that address the user's specific request."""
messages = [{"role": "system", "content": system_message}]
messages.extend(history)
messages.append({"role": "user", "content": message})
response = ""
for message_chunk in client.chat_completion(
messages,
max_tokens=512,
stream=True,
temperature=0.7,
top_p=0.9,
):
choices = message_chunk.choices
token = ""
if len(choices) and choices[0].delta.content:
token = choices[0].delta.content
response += token
yield response
def parse_ai_structures(ai_response, selected_smiles):
"""Parse AI response to extract SMILES strings and generate images."""
structures = []
lines = ai_response.split('\n')
for line in lines:
if 'STRUCTURE' in line.upper() and ':' in line:
try:
# Extract SMILES from lines like "STRUCTURE 1: CCO - ethanol"
parts = line.split(':', 1)
if len(parts) > 1:
smiles_part = parts[1].split('-')[0].strip()
# Clean up the SMILES string
smiles = smiles_part.strip('[]()').strip()
# Validate and generate image
mol = Chem.MolFromSmiles(smiles)
if mol:
img = Draw.MolToImage(mol, size=(150, 150), kekulize=True)
structures.append((img, f"Generated: {smiles}"))
except:
continue
# If no structures were parsed, generate some variations of the selected molecule
if not structures and selected_smiles:
try:
mol = Chem.MolFromSmiles(selected_smiles)
if mol:
# Generate a few variations with different sizes
for size in [(120, 120), (150, 150), (180, 180)]:
img = Draw.MolToImage(mol, size=size, kekulize=True)
structures.append((img, f"Variation: {selected_smiles}"))
except:
pass
return structures
def handle_structure_chat(message, history, selected_smiles, hf_token: str):
"""Handle the structure generation chat."""
if not message.strip():
return history, []
# Add user message to history
history.append({"role": "user", "content": message})
# Generate AI response
ai_response = ""
for chunk in generate_ai_structures(message, history[:-1], selected_smiles, hf_token):
ai_response = chunk
# Add AI response to history
history.append({"role": "assistant", "content": ai_response})
# Parse and generate structure images
structures = parse_ai_structures(ai_response, selected_smiles)
return history, structures
|