update initialize_model_once, create_llm_pipeline, force cpu only

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import os
+os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Force CPU only
 import uuid
 import threading
 import pandas as pd
@@ -88,26 +89,37 @@ def initialize_model_once(model_key):
             model_info = MODEL_CONFIG[model_key]
             model_name = model_info["name"]
             MODEL_CACHE["model_name"] = model_key
-            
-            # Handle T5 models separately
-            if model_info.get("is_t5", False):
-                MODEL_CACHE["tokenizer"] = T5Tokenizer.from_pretrained(model_name)
-                MODEL_CACHE["model"] = T5ForConditionalGeneration.from_pretrained(
-                    model_name,
-                    torch_dtype=model_info["dtype"],
-                    device_map="auto",
-                    low_cpu_mem_usage=True
-                )
-            else:
-                # Load tokenizer and model with appropriate configuration
-                MODEL_CACHE["tokenizer"] = AutoTokenizer.from_pretrained(model_name)
-                MODEL_CACHE["model"] = AutoModelForCausalLM.from_pretrained(
-                    model_name,
-                    torch_dtype=model_info["dtype"],
-                    device_map="auto",
-                    low_cpu_mem_usage=True,
-                    trust_remote_code=True
-                )
+
+            try:
+                print(f"Loading model: {model_name}")
+                # Handle T5 models separately
+                if model_info.get("is_t5", False):
+                    MODEL_CACHE["tokenizer"] = T5Tokenizer.from_pretrained(model_name)
+                    MODEL_CACHE["model"] = T5ForConditionalGeneration.from_pretrained(
+                        model_name,
+                        torch_dtype=model_info["dtype"],
+                        device_map="auto" if torch.cuda.is_available() else None,
+                        low_cpu_mem_usage=True
+                    )
+                else:
+                    # Load tokenizer and model with appropriate configuration
+                    MODEL_CACHE["tokenizer"] = AutoTokenizer.from_pretrained(model_name)
+                    MODEL_CACHE["model"] = AutoModelForCausalLM.from_pretrained(
+                        model_name,
+                        torch_dtype=model_info["dtype"],
+                        device_map="auto" if torch.cuda.is_available() else None,
+                        low_cpu_mem_usage=True,
+                        trust_remote_code=True
+                    )
+                print(f"Model {model_name} loaded successfully")
+            except Exception as e:
+                import traceback
+                print(f"Error loading model {model_name}: {str(e)}")
+                print(traceback.format_exc())
+                raise RuntimeError(f"Failed to load model {model_name}: {str(e)}")
+                
+    if MODEL_CACHE["model"] is None or MODEL_CACHE["tokenizer"] is None:
+        raise ValueError(f"Model or tokenizer not initialized properly for {model_key}")            
 
     return MODEL_CACHE["tokenizer"], MODEL_CACHE["model"], model_info.get("is_t5", False)
 
@@ -117,6 +129,9 @@ def create_llm_pipeline(model_key):
         print(f"Creating pipeline for model: {model_key}")
         tokenizer, model, is_t5 = initialize_model_once(model_key)
 
+        if model is None or tokenizer is None:
+            raise ValueError(f"Model or tokenizer is None for {model_key}")
+
         # Create appropriate pipeline based on model type
         if is_t5:
             print("Creating T5 pipeline")
@@ -124,7 +139,7 @@ def create_llm_pipeline(model_key):
                 "text2text-generation",
                 model=model,
                 tokenizer=tokenizer,
-                max_new_tokens=256,
+                max_new_tokens=128,  # Reduced for better performance
                 temperature=0.3,
                 top_p=0.9,
                 return_full_text=False,
@@ -135,7 +150,7 @@ def create_llm_pipeline(model_key):
                 "text-generation",
                 model=model,
                 tokenizer=tokenizer,
-                max_new_tokens=256,
+                max_new_tokens=128,  # Reduced for better performance
                 temperature=0.3,
                 top_p=0.9,
                 top_k=30,
@@ -150,7 +165,7 @@ def create_llm_pipeline(model_key):
         import traceback
         print(f"Error creating pipeline: {str(e)}")
         print(traceback.format_exc())
-        raise
+        raise RuntimeError(f"Failed to create pipeline: {str(e)}")
         
 def create_conversational_chain(db, file_path, model_key):
     llm = create_llm_pipeline(model_key)