add modal backend

modal-medgemma.py

@@ -0,0 +1,269 @@
+"""
+modal-medgemma.py
+-----------------
+
+This module defines a Modal app and FastAPI endpoint for running the MedGemma agent, a multimodal LLM that can process text and images. It provides a streaming API for inference, including Wikipedia tool-calling capabilities, and handles model download, loading, and inference with GPU support.
+
+Nb. needs to be deployed with the following command:
+`modal deploy modal-medgemma.py`
+
+Key components:
+- Modal app and volume setup for model weights
+- MedGemmaAgent class for model loading and inference
+- FastAPI endpoint for streaming responses
+- Utility for processing base64-encoded images
+"""
+import modal
+from typing import Optional, Generator, Dict, Any, List
+import os
+from fastapi import Security, HTTPException, Depends
+from fastapi.security.api_key import APIKeyHeader
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+import json
+import base64
+from PIL import Image
+import io
+
+app = modal.App("example-medgemma-agent")
+volume = modal.Volume.from_name("model-weights-vol", create_if_missing=True)
+MODEL_DIR = "/models"
+# MODEL_ID = "google/medgemma-4b-it"
+MODEL_ID = "unsloth/medgemma-4b-it-unsloth-bnb-4bit"
+MINUTES = 60
+
+API_KEY_NAME = "X-API-Key"
+api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=True)
+
+async def get_api_key(api_key_header: str = Security(api_key_header)):
+    """
+    Validates the provided API key against the environment variable.
+
+    Args:
+        api_key_header (str): The API key provided in the request header.
+
+    Raises:
+        HTTPException: If the API key is invalid.
+
+    Returns:
+        str: The validated API key.
+    """
+    if api_key_header != os.environ["FASTAPI_KEY"]:
+        raise HTTPException(
+            status_code=403, detail="Invalid API Key"
+        )
+    return api_key_header
+
+image = (
+    modal.Image.debian_slim()
+    .pip_install(
+        "smolagents[vllm]",
+        "fastapi[standard]",
+        "wikipedia-api",
+        "accelerate",
+        "bitsandbytes",
+        "huggingface-hub[hf_transfer]",
+        "Pillow")
+    .env({
+        "HF_HUB_ENABLE_HF_TRANSFER": "1", 
+        "HF_HUB_CACHE": MODEL_DIR
+    })
+)
+
+with image.imports():
+    from smolagents import VLLMModel, ToolCallingAgent, tool
+    from pydantic import BaseModel
+    import wikipediaapi
+
+@app.function(
+    image=image,
+    secrets=[modal.Secret.from_name("access_medgemma_hf")],
+    volumes={MODEL_DIR: volume}
+)
+def download_model():
+    """
+    Downloads the model weights from Hugging Face Hub using the provided token.
+
+    Returns:
+        dict: Status message indicating success.
+    """
+    from huggingface_hub import snapshot_download
+    snapshot_download(
+        repo_id=MODEL_ID,
+        token=os.environ["HF_TOKEN"]
+    )
+    return {"status": "Model downloaded successfully"}
+
+@app.cls(
+    image=image,
+    gpu="L4:1",
+    volumes={MODEL_DIR: volume},
+    min_containers=1,
+    max_containers=1,
+    timeout=15 * MINUTES,
+    secrets=[modal.Secret.from_name("access_medgemma_hf")],
+)
+class MedGemmaAgent:
+    """
+    Modal class for managing the MedGemma model and running inference with optional tool-calling.
+    """
+    @modal.enter()
+    def load_models(self):
+        """
+        Loads the MedGemma model into memory and prepares it for inference.
+        Downloads the model weights if not already present.
+        """
+        download_model.remote()
+        model_kwargs = {
+            "max_model_len": 8192,
+            "dtype": "bfloat16",
+            "gpu_memory_utilization": 0.95,
+            "tensor_parallel_size": 1,
+            "trust_remote_code": True
+        }
+        self.model = VLLMModel(
+            model_id=MODEL_ID,
+            model_kwargs=model_kwargs
+        )
+        print(f"Model: {MODEL_ID} loaded successfully")
+
+    @modal.method()
+    def run(self, prompt: str, images: Optional[List[Image.Image]] = None) -> Generator[Dict[str, Any], None, None]:
+        """
+        Runs the MedGemma agent on the provided prompt and optional images, yielding streaming responses.
+
+        Args:
+            prompt (str): The user prompt to process.
+            images (Optional[List[Image.Image]]): List of PIL Images to provide as context (optional).
+
+        Yields:
+            Dict[str, Any]: Streaming response chunks, including 'thinking' and 'final' messages.
+        """
+        @tool
+        def wiki(query: str) -> str:
+            """
+            Fetches a summary of a Wikipedia page based on a given search query (only one word or group of words).
+
+            Args:
+                query: The search term for the Wikipedia page (only one word or group of words).
+            """
+            wiki = wikipediaapi.Wikipedia(language="en", user_agent="MinimalAgent/1.0")
+            page = wiki.page(query)
+            if not page.exists():
+                return "No Wikipedia page found."
+            return page.summary[:1000]
+
+        self.agent = ToolCallingAgent(
+            tools=[wiki],
+            model=self.model,
+            max_steps=3
+        )
+
+        # Yield thinking step
+        yield {
+            "type": "thinking",
+            "content": {"message": "Starting to process your request..."}
+        }
+
+        # Run the agent and capture the result
+        result = self.agent.run(
+            task=prompt,
+            stream=False,
+            reset=True,
+            images=images if images else None,
+            additional_args={"flatten_messages_as_text": False},
+            max_steps=3
+        )
+        
+        # Yield the final response
+        yield {
+            "type": "final",
+            "content": {"response": result}
+        }
+
+class PromptRequest(BaseModel):
+    """
+    Request model for the /run_medgemma endpoint.
+
+    Attributes:
+        prompt (str): The user prompt to process.
+        image (Optional[str]): Base64-encoded image string (optional).
+        history (Optional[list]): Conversation history (optional).
+    """
+    prompt: str
+    image: Optional[str] = None  # Base64 encoded image
+    history: Optional[list] = None
+
+class GenerationResponse(BaseModel):
+    """
+    Response model for non-streaming generation (not used in this file).
+
+    Attributes:
+        response (str): The generated response.
+    """
+    response: str
+
+class StreamResponse(BaseModel):
+    """
+    Response model for streaming responses.
+
+    Attributes:
+        type (str): The type of message ('thinking', 'tool_call', 'tool_result', 'final').
+        content (Dict[str, Any]): The content of the message.
+    """
+    type: str  # 'thinking', 'tool_call', 'tool_result', 'final'
+    content: Dict[str, Any]
+
+def process_image(image_base64: Optional[str]) -> Optional[Image.Image]:
+    """
+    Decodes a base64-encoded image string into a PIL Image.
+
+    Args:
+        image_base64 (Optional[str]): Base64-encoded image string.
+
+    Returns:
+        Optional[Image.Image]: The decoded PIL Image, or None if decoding fails or input is None.
+    """
+    if not image_base64:
+        return None
+    try:
+        image_data = base64.b64decode(image_base64)
+        return Image.open(io.BytesIO(image_data))
+    except Exception as e:
+        print(f"Error processing image: {e}")
+        return None
+
+@app.function(
+    image=image,
+    secrets=[
+        modal.Secret.from_name("access_medgemma_hf"),
+        modal.Secret.from_name("FASTAPI_KEY")
+    ]
+)
+@modal.fastapi_endpoint(method="POST")
+async def run_medgemma(request: PromptRequest, api_key: str = Depends(get_api_key)):
+    """
+    FastAPI endpoint for running the MedGemma agent with streaming responses.
+
+    Args:
+        request (PromptRequest): The request payload containing prompt and optional image.
+        api_key (str): The validated API key (injected by Depends).
+
+    Returns:
+        StreamingResponse: An event-stream response yielding model output chunks.
+    """
+    model_handler = MedGemmaAgent()
+    
+    # Process image if provided
+    image = process_image(request.image)
+    images = [image] if image else None
+    
+    async def generate():
+        stream = model_handler.run.remote_gen.aio(request.prompt, images=images)
+        async for chunk in stream:
+            yield f"data: {json.dumps(chunk)}\n\n"
+    
+    return StreamingResponse(
+        generate(),
+        media_type="text/event-stream"
+    )