api.py

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)}")