Update app.py

app.py

@@ -1,7 +1,6 @@
 import spaces
 import torch
 import gradio as gr
-import re
 from transformers import AutoProcessor, AutoModelForImageTextToText
 MODEL_ID = "google/medgemma-1.5-4b-it"
 processor = AutoProcessor.from_pretrained(MODEL_ID)
@@ -10,76 +9,12 @@ UNUSED95_ID = processor.tokenizer.convert_tokens_to_ids('<unused95>')
 EOT_ID = processor.tokenizer.convert_tokens_to_ids('<end_of_turn>')
 def extract_response(output_ids, input_length):
     ids = output_ids.tolist()
-    
-    # Strategy 1: MedGemma uses <unused95> to separate thought from answer
     if UNUSED95_ID in ids:
         idx = ids.index(UNUSED95_ID)
         response_ids = ids[idx + 1:]
     else:
-        # Strategy 2: Skip the input prompt and look for answer in the rest
         response_ids = ids[input_length:]
-    
-    text = processor.decode(response_ids, skip_special_tokens=True).strip()
-    
-    # Strip explicit "thought" prefix if present
-    text = re.sub(r'^thought\s*\n?', '', text, flags=re.IGNORECASE)
-    
-    # Strategy 3: If the output is mixed reasoning + answer, find where Spanish starts
-    # Look for the first substantial Spanish sentence
-    spanish_markers = [
-        r'(?:^|\n{2,})(El paciente\s)',
-        r'(?:^|\n{2,})(Se observa\s)',
-        r'(?:^|\n{2,})(Los resultados\s)',
-        r'(?:^|\n{2,})(La interpretación\s)',
-        r'(?:^|\n{2,})(Hallazgos\s)',
-        r'(?:^|\n{2,})(Recomendaciones\s)',
-        r'(?:^|\n{2,})(Análisis\s)',
-        r'(?:^|\n{2,})(Interpretación\s)',
-        r'(?:^|\n{2,})(En resumen\s)',
-        r'(?:^|\n{2,})(Conclusiones\s)',
-        r'(?:^|\n{2,})(Diagnóstico\s)',
-        r'(?:^|\n{2,})(Evaluación\s)',
-        r'(?:^|\n{2,})(Nota\s)',
-        r'(?:^|\n{2,})(Observaciones\s)',
-        r'(?:^|\n{2,})(Resumen\s)',
-        r'(?:^|\n{2,})(Síntesis\s)',
-        r'(?:^|\n{2,})(Conclusión\s)',
-        r'(?:^|\n{2,})([ÁÉÍÓÚÑ][a-záéíóúñ]+)',  # Any Spanish sentence starting with capital accented letter
-    ]
-    
-    for pattern in spanish_markers:
-        match = re.search(pattern, text, re.IGNORECASE)
-        if match:
-            start_idx = match.start()
-            if start_idx > 0:
-                text = text[start_idx:]
-            break
-    
-    # Strip remaining reasoning artifacts
-    text = re.sub(r"Here'?s a thinking process[\s\S]*?(?=\n{2,}[A-ZÁÉÍÓÚÑ]|\Z)", '', text, flags=re.IGNORECASE)
-    text = re.sub(r"Understand the Role[\s\S]*?(?=\n{2,}[A-ZÁÉÍÓÚÑ]|\Z)", '', text, flags=re.IGNORECASE)
-    text = re.sub(r"Analyze the Request[\s\S]*?(?=\n{2,}[A-ZÁÉÍÓÚÑ]|\Z)", '', text, flags=re.IGNORECASE)
-    text = re.sub(r"Review the Lab Results[\s\S]*?(?=\n{2,}[A-ZÁÉÍÓÚÑ]|\Z)", '', text, flags=re.IGNORECASE)
-    text = re.sub(r"Synthesize Findings[\s\S]*?(?=\n{2,}[A-ZÁÉÍÓÚÑ]|\Z)", '', text, flags=re.IGNORECASE)
-    text = re.sub(r"Formulate Clinical Interpretation[\s\S]*?(?=\n{2,}[A-ZÁÉÍÓÚÑ]|\Z)", '', text, flags=re.IGNORECASE)
-    text = re.sub(r"Refine Interpretation[\s\S]*?(?=\n{2,}[A-ZÁÉÍÓÚÑ]|\Z)", '', text, flags=re.IGNORECASE)
-    
-    # Strip numbered step headers
-    text = re.sub(r'^\d+\.\s*\*\*[^*]+\*\*[\s\S]*?(?=\n{2,}[A-ZÁÉÍÓÚÑ]|\Z)', '', text, flags=re.MULTILINE)
-    text = re.sub(r'^\d+\.\s*\*\*[^*]+\*\*\s*$', '', text, flags=re.MULTILINE)
-    
-    # Final cleanup
-    text = text.strip()
-    text = re.sub(r'\n{3,}', '\n\n', text)
-    
-    # Strategy 4: Validate we have actual clinical content
-    has_spanish = bool(re.search(r'[áéíóúñÁÉÍÓÚÑ]', text))
-    has_substance = len(text) > 50  # At least 50 chars of actual text
-    
-    if not has_spanish or not has_substance:
-        return "ERROR: El modelo generó razonamiento interno pero no alcanzó a producir la interpretación clínica dentro del límite de tokens. Intenta con un análisis más breve o sin imágenes."
-    
-    return text
+    return processor.decode(response_ids, skip_special_tokens=True).strip()
 @spaces.GPU(duration=90)
 def analyze(image1, image2, image3, image4, text_prompt):
     images = [img for img in [image1, image2, image3, image4] if img is not None]
@@ -91,15 +26,8 @@ def analyze(image1, image2, image3, image4, text_prompt):
         messages, add_generation_prompt=True, tokenize=True,
         return_dict=True, return_tensors="pt"
     ).to(model.device, dtype=torch.bfloat16)
-    
     with torch.inference_mode():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=2048,
-            do_sample=False,
-            temperature=0.1,
-        )
-    
+        output = model.generate(**inputs, max_new_tokens=2048, eos_token_id=EOT_ID)
     return extract_response(output[0], inputs['input_ids'].shape[1])
 demo = gr.Interface(
     fn=analyze,