#!/usr/bin/env python3 import os import torch import gc import logging from typing import Dict from pydantic import BaseModel # ML Libraries from transformers import ( AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM, BitsAndBytesConfig, Gemma3ForConditionalGeneration ) from huggingface_hub import login from peft import PeftModel import warnings # API Framework from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware warnings.filterwarnings("ignore") logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) class GenerationRequest(BaseModel): input_text: str model_name: str class GenerationResponse(BaseModel): response: str class MedicalKnowledgeTester: def __init__(self): self.device = "cuda" if torch.cuda.is_available() else "cpu" logger.info(f"Using device: {self.device}") hf_token = os.environ.get("HF_TOKEN") if hf_token: login(token=hf_token) logger.info("Logged in to Hugging Face using token from environment variable.") self.models = {} self.tokenizers = {} # --- MODIFICATION: Added 'med_gemma' to the model configurations --- # NOTE: I have assumed 'google/med-gemma-2b' as the base model and a corresponding adapter name. # Please verify these are the correct Hugging Face IDs for your model. self.model_configs = { "led-base": { "model_type": "encoder-decoder", "base_model": "allenai/led-base-16384", "adapter_model": "ALQAMARI/led-base-sbar-summary-adapter", "max_length": 4096, "use_quantization": False, }, "gemma-3-12b-it": { "model_type": "decoder", "base_model": "google/gemma-3-12b-it", "adapter_model": "ALQAMARI/gemma-3-12b-it-summary-adapter", "max_length": 4096, "use_quantization": True, }, # --- CORRECTED: Using the new 27B model and adapter --- "medgemma-27b": { "model_type": "decoder", "base_model": "google/medgemma-27b-text-it", # Corrected base model "adapter_model": "ALQAMARI/medgemma-sbar-summary", # Corrected adapter [cite: 2] "max_length": 4096, "use_quantization": True, # Quantization is essential for this model size } } self.GENERAL_TEMPLATE = """You are a versatile and highly skilled medical AI assistant. Your role is to provide accurate and helpful responses to medical inquiries. - If the user provides a patient record, a long medical report, or text that requires summarization, your primary task is to summarize it concisely. Highlight the key findings, diagnoses, and recommendations in a clear format suitable for other medical professionals. - If the user asks a direct question, provide a comprehensive and clear medical explanation. - Analyze the user's input below and respond in the most appropriate manner, either as a summarizer or a knowledge expert. User Input: {input_text} Your Response:""" def load_model(self, model_name: str): if model_name in self.models: logger.info(f"Model '{model_name}' is already loaded.") return if model_name not in self.model_configs: raise ValueError(f"Model {model_name} not supported.") config = self.model_configs[model_name] logger.info(f"Loading {model_name}...") model_kwargs = {"device_map": "auto", "trust_remote_code": True} if config["use_quantization"]: bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True ) model_kwargs["quantization_config"] = bnb_config model_kwargs["torch_dtype"] = torch.bfloat16 else: model_kwargs["torch_dtype"] = torch.float16 tokenizer = AutoTokenizer.from_pretrained(config["base_model"]) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "left" if config["model_type"] == "encoder-decoder": base_model = AutoModelForSeq2SeqLM.from_pretrained(config["base_model"], **model_kwargs) else: base_model = AutoModelForCausalLM.from_pretrained(config["base_model"], **model_kwargs) try: model = PeftModel.from_pretrained(base_model, config["adapter_model"]) logger.info(f"Successfully loaded adapter from {config['adapter_model']}") except Exception as e: logger.error(f"Failed to load adapter: {e}. Using base model without adapter.") model = base_model model.eval() self.models[model_name] = model self.tokenizers[model_name] = tokenizer logger.info(f"{model_name} loaded successfully.") def generate_response(self, model_name: str, input_text: str) -> str: if model_name not in self.models: self.load_model(model_name) model = self.models[model_name] tokenizer = self.tokenizers[model_name] config = self.model_configs[model_name] prompt = self.GENERAL_TEMPLATE.format(input_text=input_text) if config["model_type"] == "decoder": inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=config["max_length"]).to(self.device) else: # For encoder-decoder models, we might just send the raw text inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=config["max_length"]).to(self.device) with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=512, do_sample=True, temperature=0.1, pad_token_id=tokenizer.eos_token_id, repetition_penalty=1.1 ) if config["model_type"] == "decoder": input_length = inputs.input_ids.shape[1] generated_tokens = outputs[0][input_length:] else: generated_tokens = outputs[0] response = tokenizer.decode(generated_tokens, skip_special_tokens=True) return response.strip() app = FastAPI() tester = MedicalKnowledgeTester() app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) @app.on_event("startup") async def startup_event(): logger.info("Server starting up. Pre-loading default model...") try: # You might want to change the default pre-loaded model or pre-load all of them tester.load_model("medgemma-27b") except Exception as e: logger.error(f"Could not pre-load medgemma-27b model: {e}s") @app.get("/") def read_root(): return {"status": "Medical AI API - I AM THE NEW VERSION"} @app.post("/generate", response_model=GenerationResponse) async def generate(request: GenerationRequest): logger.info(f"Received request for model: {request.model_name}") try: response_text = tester.generate_response( model_name=request.model_name, input_text=request.input_text ) return GenerationResponse(response=response_text) except Exception as e: logger.error(f"Error during generation: {e}") raise HTTPException(status_code=500, detail=str(e))