Update app.py

app.py

@@ -24,8 +24,6 @@ warnings.filterwarnings("ignore")
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# --- MODIFIED: Simplified API Request Model ---
-# We no longer need 'task_type' from the frontend.
 class GenerationRequest(BaseModel):
     input_text: str
     model_name: str
@@ -46,6 +44,9 @@ class MedicalKnowledgeTester:
         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",
@@ -60,11 +61,16 @@ class MedicalKnowledgeTester:
                 "adapter_model": "ALQAMARI/gemma-3-12b-it-summary-adapter",
                 "max_length": 4096,
                 "use_quantization": True,
+            },
+            "med_gemma": {
+                "model_type": "decoder",
+                "base_model": "google/med-gemma-2b", # Assumed base model, please verify
+                "adapter_model": "ALQAMARI/med-gemma-summary-adapter", # Assumed adapter, please verify
+                "max_length": 4096,
+                "use_quantization": True, # Assumed quantization for efficiency
             }
         }
         
-        # --- MODIFIED: New General-Purpose Prompt Template ---
-        # This single prompt instructs the model to be smart about the user's input.
         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.
@@ -121,8 +127,6 @@ Your Response:"""
         self.tokenizers[model_name] = tokenizer
         logger.info(f"{model_name} loaded successfully.")
 
-    # --- MODIFIED: Simplified generate_response function ---
-    # It no longer needs 'task_type' and uses the general template for everything.
     def generate_response(self, model_name: str, input_text: str) -> str:
         if model_name not in self.models:
             self.load_model(model_name)
@@ -168,6 +172,7 @@ app.add_middleware(
 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("gemma-3-12b-it")
     except Exception as e:
         logger.error(f"Could not pre-load gemma-3-12b-it model: {e}")
@@ -176,7 +181,6 @@ async def startup_event():
 def read_root():
     return {"status": "Medical AI API - I AM THE NEW VERSION"}
 
-# --- MODIFIED: Updated generate endpoint ---
 @app.post("/generate", response_model=GenerationResponse)
 async def generate(request: GenerationRequest):
     logger.info(f"Received request for model: {request.model_name}")