Update app.py

app.py

@@ -46,6 +46,9 @@ class ModelManager:
     def initialize_models(self):
         """Initialize models with optimized settings"""
         try:
+            import torch
+            from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+            
             HUGGINGFACE_TOKEN = os.environ.get('HUGGINGFACE_TOKEN')
             if not HUGGINGFACE_TOKEN:
                 raise ValueError("HUGGINGFACE_TOKEN environment variable not set")
@@ -53,6 +56,14 @@ class ModelManager:
             logger.info("Starting model initialization...")
             model_name = "meta-llama/Llama-2-7b-chat-hf"
 
+            # Configure 4-bit quantization
+            bnb_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=torch.bfloat16
+            )
+
             # Load tokenizer with optimized settings
             logger.info("Loading tokenizer...")
             self.tokenizer = AutoTokenizer.from_pretrained(
@@ -63,43 +74,25 @@ class ModelManager:
             )
             self.tokenizer.pad_token = self.tokenizer.eos_token
 
-            # Initialize model with Unsloth optimizations
-            logger.info("Loading model with Unsloth optimizations...")
-            model, tokenizer = FastLanguageModel.from_pretrained(
-                model_name=model_name,
+            # Initialize model with optimized settings
+            logger.info("Loading model...")
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_name,
                 token=HUGGINGFACE_TOKEN,
-                max_seq_length=512,
-                dtype="float16",
-                load_in_4bit=True,  # Use 4-bit quantization
-                device_map="auto",  # Automatically handle device mapping
-                kwargs=dict(
-                    use_gradient_checkpointing=True,
-                    use_flash_attention_2=True,
-                    use_merged_kernels=True,
-                )
-            )
-            
-            # Apply additional optimizations
-            model = FastLanguageModel.get_peft_model(
-                model,
-                r=16,
-                target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
-                modules_to_save=None,
-                lora_alpha=16,
-                lora_dropout=0.05,
-                bias="none",
-                use_gradient_checkpointing=True,
-                random_state=42,
-                use_rslora=False,
-                use_dora=False,
+                device_map="auto",
+                torch_dtype=torch.float16,
+                quantization_config=bnb_config,
+                use_flash_attention_2=True,
+                use_cache=True,
+                attn_implementation="flash_attention_2",
+                low_cpu_mem_usage=True,
             )
 
-            self.model = model
-            logger.info("Model loaded successfully with Unsloth optimizations")
-
             # Create optimized pipeline
             logger.info("Creating pipeline...")
-            self.news_generator = FastLanguageModel.get_pipeline(
+            from transformers import pipeline
+            self.news_generator = pipeline(
+                "text-generation",
                 model=self.model,
                 tokenizer=self.tokenizer,
                 device_map="auto",