Update app.py

app.py

@@ -47,7 +47,7 @@ class ModelManager:
         """Initialize models with optimized settings"""
         try:
             import torch
-            from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+            from transformers import AutoModelForCausalLM, AutoTokenizer
             
             HUGGINGFACE_TOKEN = os.environ.get('HUGGINGFACE_TOKEN')
             if not HUGGINGFACE_TOKEN:
@@ -56,14 +56,6 @@ class ModelManager:
             logger.info("Starting model initialization...")
             model_name = "meta-llama/Llama-2-7b-chat-hf"
 
-            # Configure 8-bit quantization instead of 4-bit
-            bnb_config = BitsAndBytesConfig(
-                load_in_8bit=True,
-                bnb_8bit_use_double_quant=True,
-                bnb_8bit_quant_type="nf8",
-                bnb_8bit_compute_dtype=torch.float16
-            )
-
             # Load tokenizer with optimized settings
             logger.info("Loading tokenizer...")
             self.tokenizer = AutoTokenizer.from_pretrained(
@@ -74,18 +66,18 @@ class ModelManager:
             )
             self.tokenizer.pad_token = self.tokenizer.eos_token
 
-            # Initialize model with optimized settings
+            # Initialize model with basic settings
             logger.info("Loading model...")
             self.model = AutoModelForCausalLM.from_pretrained(
                 model_name,
                 token=HUGGINGFACE_TOKEN,
                 device_map="auto",
                 torch_dtype=torch.float16,
-                quantization_config=bnb_config,
+                load_in_8bit=True,
                 low_cpu_mem_usage=True,
             )
 
-            # Create optimized pipeline
+            # Create pipeline
             logger.info("Creating pipeline...")
             from transformers import pipeline
             self.news_generator = pipeline(
@@ -103,11 +95,11 @@ class ModelManager:
                 early_stopping=True
             )
 
-            # Load Whisper model with optimized settings
+            # Load Whisper model with basic settings
             logger.info("Loading Whisper model...")
             self.whisper_model = whisper.load_model(
                 "tiny",
-                device="cuda",
+                device="cuda" if torch.cuda.is_available() else "cpu",
                 download_root="/tmp/whisper",
                 in_memory=True
             )