update.py

app.py

@@ -8,41 +8,34 @@ logging.basicConfig(level=logging.DEBUG,
                     format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Models to try
-MODELS_TO_TRY = [
-    "google/flan-t5-xxl",  # Powerful instruction-following model
-    "bigscience/T0pp",     # Optimized for zero-shot tasks
-    "t5-large",            # General-purpose text generation
-    "google/flan-t5-large" # Lightweight instruction-tuned model
-]
+# Model to use
+MODEL_NAME = "google/flan-t5-large"
 
 def load_model():
     """
-    Attempt to load a suitable model for text generation.
+    Load the selected model and tokenizer using PyTorch.
     """
-    for model_name in MODELS_TO_TRY:
-        try:
-            logger.info(f"Attempting to load model: {model_name}")
-            
-            # Load model and tokenizer
-            model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            
-            # Create the text generation pipeline
-            generator = pipeline(
-                "text2text-generation",
-                model=model,
-                tokenizer=tokenizer,
-                max_length=512,
-                num_return_sequences=1
-            )
-            logger.info(f"Successfully loaded model: {model_name}")
-            return generator
-        except Exception as e:
-            logger.error(f"Failed to load model {model_name}: {e}")
-    
-    logger.error("All model attempts failed. No model loaded.")
-    return None
+    try:
+        logger.info(f"Loading model: {MODEL_NAME} with PyTorch backend")
+        
+        # Load the model and tokenizer
+        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, framework="pt")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        
+        # Create the text generation pipeline
+        generator = pipeline(
+            "text2text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            framework="pt",  # Specify PyTorch framework
+            max_length=512,
+            num_return_sequences=1
+        )
+        logger.info(f"Successfully loaded model: {MODEL_NAME}")
+        return generator
+    except Exception as e:
+        logger.error(f"Failed to load model {MODEL_NAME}: {e}")
+        return None
 
 # Load the generator at startup
 generator = load_model()