import os import sys import asyncio import numpy as np import pandas as pd import io import time import logging from datetime import datetime, date from typing import Dict, List from functools import partial # Import partial for cleaner dependency definitions # --- FastAPI & Security Imports --- from fastapi import FastAPI, HTTPException, Security, Depends, Request, File, UploadFile from fastapi.security import APIKeyHeader from fastapi.responses import StreamingResponse from pydantic import BaseModel # 2. Imports for the ML Model (kept from original script) from smi_ssed.load import load_smi_ssed import torch from torch import nn from torch.utils.data import Dataset, DataLoader from tqdm import tqdm logging.basicConfig(level=logging.INFO) app = FastAPI() # FIX: Changed to an absolute path (/workspace/api_keys.db) to avoid file access issues # caused by the preceding os.chdir() call and to ensure path consistency. # API Key header api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False) # --- FastAPI Event Handlers and Routes --- # Pydantic model for response class Message(BaseModel): message: str @app.post("/refresh-keys", response_model=Message) async def refresh_keys(): # This endpoint is now purely illustrative, as keys are loaded on demand. return {"message": "API key validation is now handled via database on demand."} # Pydantic model for single SMILES input class SmileInput(BaseModel): smiles: str # --- ML Model Initialization (kept from original script) --- # This section remains largely unchanged, ensuring the model is loaded once at startup. SMILES_COLUMN_NAME = 'SMILES' FINETUNED_CHECKPOINT_PATH = '/app/CDI/model/finetuned_model.pth' YOUR_EMBEDDING_DIM = 8192 MAX_LENGTH = 128 BATCH_SIZE = 32 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using device: {device}") # --- 1. Model Definition --- class SmiSsedPredictor(nn.Module): def __init__(self, smi_ssed_model, embedding_dim): super(SmiSsedPredictor, self).__init__() self.base_model_encoder = smi_ssed_model.encoder true_hidden_dim = self.base_model_encoder.mamba.embedding.weight.shape[1] self.regressor = nn.Linear(true_hidden_dim, embedding_dim) def forward(self, input_ids, attention_mask): outputs = self.base_model_encoder(input_ids, mask=attention_mask) hidden_states = outputs[0] if hidden_states.dim() == 2: batch_size = input_ids.shape[0] if batch_size > 0: seq_len = hidden_states.shape[0] // batch_size hidden_states = hidden_states.view(batch_size, seq_len, -1) expanded_mask = attention_mask.unsqueeze(-1) if expanded_mask.shape[1] != hidden_states.shape[1]: mask_for_pooling = torch.zeros_like(hidden_states) slice_len = min(expanded_mask.shape[1], hidden_states.shape[1]) mask_for_pooling[:, :slice_len, :] = expanded_mask[:, :slice_len, :] else: mask_for_pooling = expanded_mask sum_hidden_states = torch.sum(hidden_states * mask_for_pooling, 1) sum_mask = torch.clamp(mask_for_pooling.sum(1), min=1e-9) molecular_representation = sum_hidden_states / sum_mask predicted_embedding = self.regressor(molecular_representation) return predicted_embedding # --- 2. Inference Dataset --- class InferenceDataset(Dataset): def __init__(self, smiles_list, tokenizer, max_length): self.tokenizer = tokenizer self.smiles = smiles_list self.max_length = max_length def __len__(self): return len(self.smiles) def __getitem__(self, index): smiles_str = self.smiles[index] encoding = self.tokenizer.encode_plus( smiles_str, add_special_tokens=True, max_length=self.max_length, padding='max_length', truncation=True, return_attention_mask=True, return_tensors='pt', ) input_ids = encoding['input_ids'] attention_mask = (input_ids != 0).float() return { 'input_ids': input_ids.flatten(), 'attention_mask': attention_mask.flatten() } try: smi_ssed_base_model = load_smi_ssed( folder="/workspace/materials.smi_ssed/smi_ssed/inference/smi_ssed", ckpt_filename='smi_ssed_130.pt' ) tokenizer = smi_ssed_base_model.tokenizer model = SmiSsedPredictor(smi_ssed_base_model, YOUR_EMBEDDING_DIM).to(device) print(f"Loading fine-tuned weights from: {FINETUNED_CHECKPOINT_PATH}") checkpoint = torch.load(FINETUNED_CHECKPOINT_PATH, map_location=device) model.load_state_dict(checkpoint['model_state_dict']) model.eval() except Exception as e: logging.error(f"Failed to load ML model dependencies: {e}") # In a real app, you might want to raise an exception or set a flag to disable endpoints async def featurizer_generator(file_content_buffer: io.StringIO): logging.info("Generator started with in-memory buffer...") # Ensure the uploaded file is a CSV try: inference_df = pd.read_csv(file_content_buffer) print("--- 1. Loading Model and Tokenizer ---") # Load the base smi_ssed model architecture and tokenizer # This is required to instantiate our custom SmiSsedPredictor class print("Model loaded successfully.") print("\n--- 2. Preparing Data for Inference ---") # Load the new SMILES data smiles_to_predict = inference_df[SMILES_COLUMN_NAME].tolist() print(f"Found {len(smiles_to_predict)} SMILES to process.") # Create dataset and dataloader inference_dataset = InferenceDataset(smiles_to_predict, tokenizer, MAX_LENGTH) inference_loader = DataLoader(inference_dataset, batch_size=BATCH_SIZE, shuffle=False) with torch.no_grad(): for batch in tqdm(inference_loader, desc="Predicting Embeddings"): try: input_ids = batch['input_ids'].to(device) attention_mask = batch['attention_mask'].to(device) batch_outputs = [] for i in range(input_ids.size(0)): single_output = model(input_ids[i].unsqueeze(0), attention_mask[i].unsqueeze(0)) batch_outputs.append(single_output) outputs = torch.cat(batch_outputs, dim=0) # Get model predictions # outputs = model(input_ids, attention_mask) outputs =outputs.cpu().numpy() current_batch_rows = outputs.shape[0] rows_to_pad = BATCH_SIZE - current_batch_rows if rows_to_pad > 0: DTYPE = outputs.dtype padding = np.zeros((rows_to_pad, YOUR_EMBEDDING_DIM), dtype=DTYPE) outputs = np.vstack([outputs, padding]) yield outputs.tobytes() logging.info(f"Streamed batch {i+1}") await asyncio.sleep(0.05) # Free up GPU memory after processing each batch del input_ids, attention_mask, batch_outputs, outputs torch.cuda.empty_cache() except Exception as e: logging.error(f"Error processing batch {i+1}: {e}") break except Exception as e: logging.error(f"Failed to read or process CSV buffer: {e}") logging.info("Finished streaming.") from typing import Annotated from typing import Annotated from fastapi import Depends, Request, Security # --- Refactored API Endpoints --- @app.post("/stream-features-from-csv") async def stream_features_from_csv( file: UploadFile = File(...) ): """ Endpoint for CSV stream featurization with database-backed rate limiting. """ file_content = (await file.read()).decode("utf-8") file_buffer = io.StringIO(file_content) return StreamingResponse( featurizer_generator(file_buffer), media_type="application/octet-stream" ) @app.post("/predict-single-smile", response_model=List[float]) async def predict_single_smile( smile_input: SmileInput ): """ Endpoint for single SMILES prediction with database-backed rate limiting. """ try: # Create dataset for single SMILES inference_dataset = InferenceDataset([smile_input.smiles], tokenizer, MAX_LENGTH) input_data = inference_dataset[0] input_ids = input_data['input_ids'].to(device).unsqueeze(0) attention_mask = input_data['attention_mask'].to(device).unsqueeze(0) # Get model prediction with torch.no_grad(): embedding = model(input_ids, attention_mask) embedding = embedding.cpu().numpy().flatten().tolist() return embedding except Exception as e: # Note: Rate limit check is done BEFORE this execution raise HTTPException(status_code=500, detail=f"Error processing SMILES: {str(e)}")