Fixed memory issue

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer, BitsAndBytesConfig
 import logging
 import gc
 import warnings
@@ -27,8 +27,8 @@ def load_model(model_key=None):
     if model_key is None:
         model_key = DEFAULT_MODEL
     
-    # Try to load models in order of preference
-    model_keys_to_try = [model_key, "meditron", "dialogpt_medium", "dialogpt_small"]
+    # Try to load models in order of preference - prioritize lightweight models
+    model_keys_to_try = [model_key, "flan_t5_small", "dialogpt_medium", "meditron"]
     
     for key in model_keys_to_try:
         if key not in MODEL_CONFIGS:
@@ -80,7 +80,11 @@ def load_model(model_key=None):
                     model_kwargs["device_map"] = None  # Let it use CPU naturally
             
             print("Loading model...")
-            model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
+            # Use appropriate model class based on model type
+            if "flan-t5" in model_name.lower() or "t5" in model_name.lower():
+                model = AutoModelForSeq2SeqLM.from_pretrained(model_name, **model_kwargs)
+            else:
+                model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
             
             current_model_name = model_name
             print(f"✅ Model loaded successfully: {model_name}")
@@ -113,7 +117,14 @@ def generate_response(prompt, max_tokens=None, temperature=None, top_p=None):
     top_p = top_p or GENERATION_DEFAULTS["top_p"]
     
     try:
-        full_prompt = f"{MEDICAL_SYSTEM_PROMPT}\n\nPatient/User: {prompt}\n"
+        # Format prompt based on model type
+        if "flan-t5" in current_model_name.lower() or "t5" in current_model_name.lower():
+            # T5 instruction format
+            full_prompt = f"{MEDICAL_SYSTEM_PROMPT}\n\nQuestion: {prompt}\nAnswer:"
+        else:
+            # Causal LM format
+            full_prompt = f"{MEDICAL_SYSTEM_PROMPT}\n\nPatient/User: {prompt}\n"
+        
         print(f"Full prompt: {full_prompt}")
         
         # Tokenize input with proper truncation
@@ -121,7 +132,7 @@ def generate_response(prompt, max_tokens=None, temperature=None, top_p=None):
             full_prompt, 
             return_tensors="pt", 
             truncation=True, 
-            max_length=512,  # Reduced for DialoGPT
+            max_length=512,
             padding=True
         )
         
@@ -129,16 +140,28 @@ def generate_response(prompt, max_tokens=None, temperature=None, top_p=None):
         device = next(model.parameters()).device
         inputs = {k: v.to(device) for k, v in inputs.items()}
         
-        # Generation parameters
-        generation_kwargs = {
-            "max_new_tokens": min(max_tokens, 1024),
-            "temperature": temperature,
-            "top_p": top_p,
-            "do_sample": GENERATION_DEFAULTS["do_sample"],
-            "pad_token_id": tokenizer.eos_token_id,
-            "repetition_penalty": GENERATION_DEFAULTS["repetition_penalty"],
-            "no_repeat_ngram_size": GENERATION_DEFAULTS["no_repeat_ngram_size"]
-        }
+        # Generation parameters - different for T5 vs causal models
+        if "flan-t5" in current_model_name.lower() or "t5" in current_model_name.lower():
+            # T5 seq2seq generation parameters
+            generation_kwargs = {
+                "max_new_tokens": min(max_tokens, 100),
+                "temperature": temperature,
+                "top_p": top_p,
+                "do_sample": GENERATION_DEFAULTS["do_sample"],
+                "repetition_penalty": GENERATION_DEFAULTS["repetition_penalty"],
+                "early_stopping": True
+            }
+        else:
+            # Causal LM generation parameters
+            generation_kwargs = {
+                "max_new_tokens": min(max_tokens, 1024),
+                "temperature": temperature,
+                "top_p": top_p,
+                "do_sample": GENERATION_DEFAULTS["do_sample"],
+                "pad_token_id": tokenizer.eos_token_id,
+                "repetition_penalty": GENERATION_DEFAULTS["repetition_penalty"],
+                "no_repeat_ngram_size": GENERATION_DEFAULTS["no_repeat_ngram_size"]
+            }
         
         print(f"Generating with kwargs: {generation_kwargs}")
         
@@ -153,9 +176,15 @@ def generate_response(prompt, max_tokens=None, temperature=None, top_p=None):
         generation_time = time.time() - start_time
         print(f"⏱️ Generation completed in {generation_time:.2f} seconds")
         
-        # Decode response and extract new content
-        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        response = full_response.replace(full_prompt, "").strip()
+        # Decode response - different handling for T5 vs causal models
+        if "flan-t5" in current_model_name.lower() or "t5" in current_model_name.lower():
+            # T5 generates only the answer, no need to remove prompt
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+        else:
+            # Causal models generate prompt + answer, need to remove prompt
+            full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            response = full_response.replace(full_prompt, "").strip()
+        
         print(f"Generated response: {response}")
         
         # Clean up response

config.py

@@ -16,8 +16,8 @@ MODEL_CONFIGS = {
     }
 }
 
-# Default model to use - prioritize medical model
-DEFAULT_MODEL = "meditron"
+# Default model to use - lightweight for 16GB memory limit
+DEFAULT_MODEL = "flan_t5_small"
 
 # Model loading settings (optimized for CPU)
 MODEL_SETTINGS = {