Update app.py

app.py

@@ -13,6 +13,7 @@ logger = logging.getLogger(__name__)
 
 # Model configuration
 MODEL_NAME = "wizcodes12/snaxfix-model"
+FALLBACK_MODEL = "google/flan-t5-small"  # Fallback model if main model fails
 SUPPORTED_LANGUAGES = [
     "python", "javascript", "java", "c", "cpp", "csharp", "rust", 
     "php", "html", "css", "sql"
@@ -73,16 +74,29 @@ class SyntaxFixerApp:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         logger.info(f"Using device: {self.device}")
         
-        # Load model and tokenizer
+        # Load model and tokenizer with fallback
+        self.model_name_used = None
         try:
+            logger.info(f"Attempting to load primary model: {MODEL_NAME}")
             self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
             self.model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
-            self.model.to(self.device)
-            self.model.eval()
-            logger.info("Model and tokenizer loaded successfully")
+            self.model_name_used = MODEL_NAME
+            logger.info("Primary model and tokenizer loaded successfully")
         except Exception as e:
-            logger.error(f"Failed to load model: {e}")
-            raise
+            logger.warning(f"Failed to load primary model: {e}")
+            logger.info(f"Attempting to load fallback model: {FALLBACK_MODEL}")
+            try:
+                self.tokenizer = AutoTokenizer.from_pretrained(FALLBACK_MODEL)
+                self.model = AutoModelForSeq2SeqLM.from_pretrained(FALLBACK_MODEL)
+                self.model_name_used = FALLBACK_MODEL
+                logger.info("Fallback model and tokenizer loaded successfully")
+            except Exception as fallback_error:
+                logger.error(f"Failed to load fallback model: {fallback_error}")
+                raise Exception(f"Failed to load both primary and fallback models. Primary: {e}, Fallback: {fallback_error}")
+        
+        self.model.to(self.device)
+        self.model.eval()
+        logger.info(f"Using model: {self.model_name_used}")
         
         # Initialize history
         self.history = []
@@ -96,8 +110,14 @@ class SyntaxFixerApp:
             return f"Error: Language '{language}' is not supported. Choose from: {', '.join(SUPPORTED_LANGUAGES)}"
         
         try:
-            # Prepare input
-            input_text = f"<{language.upper()}> Fix the syntax errors in this {language} code: {broken_code}"
+            # Prepare input - adjust prompt based on model being used
+            if self.model_name_used == FALLBACK_MODEL:
+                # Simplified prompt for fallback model
+                input_text = f"Fix the syntax errors in this {language} code: {broken_code}"
+            else:
+                # Original prompt for specialized model
+                input_text = f"<{language.upper()}> Fix the syntax errors in this {language} code: {broken_code}"
+            
             inputs = self.tokenizer(
                 input_text,
                 max_length=MAX_LENGTH,
@@ -163,7 +183,10 @@ def create_gradio_interface():
     
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("# SnaxFix: Advanced Syntax Error Fixer")
-        gr.Markdown("Fix syntax errors in code across multiple programming languages using the `snaxfix-model` by wizcodes12.")
+        gr.Markdown("Fix syntax errors in code across multiple programming languages using AI models.")
+        
+        # Show which model is being used
+        gr.Markdown(f"**Currently using model:** `{app.model_name_used}`")
         
         with gr.Row():
             with gr.Column(scale=2):