test

app.py

@@ -1,51 +1,66 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import logging
 import spaces
 
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # Dictionary to store loaded models and tokenizers
 loaded_models = {}
 
-# List of available models (update with your preferred models)
+# List of available models (ensure these are correct and accessible)
 models = [
     "hanzla/gemma-2b-datascience-instruct-v5",
     "hanzla/gemma-2b-datascience-instruct-v4.5"
 ]
 
 
-@spaces.GPU
 def load_all_models():
     """
     Pre-loads all models and their tokenizers into memory.
     """
     for model_name in models:
         if model_name not in loaded_models:
-            print(f"Loading model: {model_name}")
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda" if torch.cuda.is_available() else "cpu")
-            loaded_models[model_name] = (model, tokenizer)
-    return "All models loaded successfully."
-
+            try:
+                logger.info(f"Loading model: {model_name}")
+                tokenizer = AutoTokenizer.from_pretrained(model_name)
+                model = AutoModelForCausalLM.from_pretrained(model_name).to(
+                    "cuda" if torch.cuda.is_available() else "cpu")
+                loaded_models[model_name] = (model, tokenizer)
+                logger.info(f"Successfully loaded {model_name}")
+            except Exception as e:
+                logger.error(f"Failed to load model {model_name}: {e}")
 
 @spaces.GPU
 def get_model_response(model_name, message):
     """
     Generates a response from the specified model given a user message.
     """
-    model, tokenizer = loaded_models[model_name]
-    inputs = tokenizer(message, return_tensors="pt").to(model.device)
-
-    # Generate response with appropriate parameters
-    outputs = model.generate(
-        **inputs,
-        max_length=512,
-        do_sample=True,
-        top_p=0.95,
-        top_k=50
-    )
-
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    try:
+        model, tokenizer = loaded_models[model_name]
+        inputs = tokenizer(message, return_tensors="pt").to(model.device)
+
+        # Generate response with appropriate parameters
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=512,
+                do_sample=True,
+                top_p=0.95,
+                top_k=50
+            )
+
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+    except KeyError:
+        logger.error(f"Model {model_name} not found in loaded_models.")
+        return f"Error: Model {model_name} not loaded."
+    except Exception as e:
+        logger.error(f"Error generating response from {model_name}: {e}")
+        return f"Error generating response: {e}"
 
 
 def chat(message, history1, history2, model1, model2):
@@ -99,6 +114,9 @@ def clear_chat():
     return [], [], "Votes - 0, 0"
 
 
+# Pre-load all models before building the Gradio interface
+load_all_models()
+
 with gr.Blocks() as demo:
     gr.Markdown("# 🤖 Hugging Face Model Comparison Chat")
 
@@ -152,8 +170,5 @@ with gr.Blocks() as demo:
         outputs=[chatbot1, chatbot2, vote_text]
     )
 
-# Pre-load all models when the space starts
-load_all_models()
-
 if __name__ == "__main__":
     demo.launch()