Update app.py

app.py

@@ -17,9 +17,7 @@ from functools import lru_cache
 import gc
 import time
 from huggingface_hub import login
-from transformers import AutoTokenizer, BitsAndBytesConfig
-from unsloth import FastLanguageModel
-import tqdm
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 # Configure logging
 logging.basicConfig(
@@ -46,15 +44,17 @@ class ModelManager:
         if not self._initialized:
             self.tokenizer = None
             self.model = None
+            self.pipeline = None
             self.whisper_model = None
             self._initialized = True
             self.last_used = time.time()
     
     @spaces.GPU()
     def initialize_llm(self):
-        """Initialize LLM model with optimization"""
+        """Initialize LLM model with standard transformers"""
         try:
-            MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"
+            # Use small model for ZeroGPU compatibility
+            MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
             
             logger.info("Loading tokenizer...")
             self.tokenizer = AutoTokenizer.from_pretrained(
@@ -62,43 +62,30 @@ class ModelManager:
                 token=HUGGINGFACE_TOKEN,
                 use_fast=True,
             )
-            self.tokenizer.pad_token = self.tokenizer.eos_token
             
-            try:
-                # Try with unsloth first
-                logger.info("Attempting to load model with unsloth optimization...")
-                self.model, self.tokenizer = FastLanguageModel.from_pretrained(
-                    model_name=MODEL_NAME,
-                    token=HUGGINGFACE_TOKEN,
-                    load_in_8bit=True,
-                    max_seq_length=2048,
-                    device_map="auto"
-                )
-                
-                # Optimize with unsloth
-                self.model = FastLanguageModel.get_peft_model(
-                    self.model,
-                    r=8,
-                    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
-                    lora_alpha=8,
-                    bias="none"
-                )
-                
-                logger.info("Model loaded successfully with unsloth")
-                
-            except Exception as unsloth_error:
-                # Fallback to standard transformers
-                logger.warning(f"Unsloth optimization failed: {str(unsloth_error)}. Falling back to standard model.")
-                from transformers import AutoModelForCausalLM
-                
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    MODEL_NAME,
-                    token=HUGGINGFACE_TOKEN,
-                    device_map="auto",
-                    torch_dtype=torch.float16,
-                    load_in_8bit=True
-                )
-                logger.info("Model loaded with standard transformers")
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+            
+            # Basic memory settings for ZeroGPU
+            logger.info("Loading model...")
+            self.model = AutoModelForCausalLM.from_pretrained(
+                MODEL_NAME,
+                token=HUGGINGFACE_TOKEN,
+                device_map="auto",
+                torch_dtype=torch.float16,
+                low_cpu_mem_usage=True
+            )
+            
+            # Create text generation pipeline
+            logger.info("Creating pipeline...")
+            self.pipeline = pipeline(
+                "text-generation",
+                model=self.model,
+                tokenizer=self.tokenizer,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                max_length=2048
+            )
             
             logger.info("LLM initialized successfully")
             self.last_used = time.time()
@@ -128,7 +115,7 @@ class ModelManager:
 
     def check_llm_initialized(self):
         """Check if LLM is initialized and initialize if needed"""
-        if self.tokenizer is None or self.model is None:
+        if self.tokenizer is None or self.model is None or self.pipeline is None:
             logger.info("LLM not initialized, initializing...")
             self.initialize_llm()
         self.last_used = time.time()
@@ -154,11 +141,15 @@ class ModelManager:
                 if hasattr(self, 'tokenizer') and self.tokenizer is not None:
                     del self.tokenizer
                     
+                if hasattr(self, 'pipeline') and self.pipeline is not None:
+                    del self.pipeline
+                    
                 if hasattr(self, 'whisper_model') and self.whisper_model is not None:
                     del self.whisper_model
                 
                 self.tokenizer = None
                 self.model = None
+                self.pipeline = None
                 self.whisper_model = None
                 
                 if torch.cuda.is_available():
@@ -490,62 +481,45 @@ Follow these requirements:
 - Do not invent information
 - Be rigorous with the provided facts [/INST]"""
 
-        # Optimize for requested size
-        max_new_tokens = min(int(size * 2.5), 1024)  # Increased limit for better quality
-
-        # Generate response using optimized unsloth model
-        with torch.inference_mode():
-            try:
-                logger.info("Generating news article...")
-                # Check if we're using unsloth or standard model
-                is_unsloth = hasattr(model_manager.model, 'unsloth_module') if hasattr(model_manager.model, 'unsloth_module') else False
-                
-                # Prepare inputs
-                inputs = model_manager.tokenizer(
-                    prompt, 
-                    return_tensors="pt", 
-                    add_special_tokens=False
-                ).to(model_manager.model.device)
-                
-                # Generate with appropriate settings
-                outputs = model_manager.model.generate(
-                    **inputs,
-                    max_new_tokens=max_new_tokens,
-                    do_sample=True,
-                    temperature=0.7,
-                    top_p=0.95,
-                    repetition_penalty=1.2,
-                    pad_token_id=model_manager.tokenizer.eos_token_id,
-                    use_cache=True
-                )
-                
-                # Decode the generated text
-                if is_unsloth:
-                    # Unsloth specific decoding
-                    generated_text = model_manager.tokenizer.decode(
-                        outputs[0][inputs.input_ids.shape[1]:], 
-                        skip_special_tokens=True
-                    )
+        # Generate with standard pipeline
+        try:
+            logger.info("Generating news article...")
+            
+            # Set max length based on requested size
+            max_length = min(len(prompt.split()) + size * 2, 2048)
+            
+            # Generate using the pipeline
+            outputs = model_manager.pipeline(
+                prompt,
+                max_length=max_length,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+                repetition_penalty=1.2,
+                pad_token_id=model_manager.tokenizer.eos_token_id,
+                num_return_sequences=1
+            )
+            
+            # Extract generated text
+            generated_text = outputs[0]['generated_text']
+            
+            # Clean up the result by removing the prompt
+            if "[/INST]" in generated_text:
+                news_article = generated_text.split("[/INST]")[1].strip()
+            else:
+                # Try to extract the text after the prompt
+                prompt_words = prompt.split()[:50]  # Use first 50 words to identify
+                prompt_fragment = " ".join(prompt_words)
+                if prompt_fragment in generated_text:
+                    news_article = generated_text[generated_text.find(prompt_fragment) + len(prompt_fragment):].strip()
                 else:
-                    # Standard transformers decoding
-                    generated_text = model_manager.tokenizer.decode(
-                        outputs[0], 
-                        skip_special_tokens=True
-                    )
-                    # Remove the prompt from the generated text
-                    prompt_text = model_manager.tokenizer.decode(
-                        inputs.input_ids[0], 
-                        skip_special_tokens=True
-                    )
-                    generated_text = generated_text.replace(prompt_text, "")
-                
-                # Clean up the generated text
-                news_article = generated_text.strip()
-                logger.info(f"News generation completed: {len(news_article)} chars")
-                
-            except Exception as gen_error:
-                logger.error(f"Error in text generation: {str(gen_error)}")
-                raise
+                    news_article = generated_text
+            
+            logger.info(f"News generation completed: {len(news_article)} chars")
+            
+        except Exception as gen_error:
+            logger.error(f"Error in text generation: {str(gen_error)}")
+            raise
         
         return news_article, raw_transcriptions
 
@@ -710,12 +684,11 @@ def create_demo():
     return demo
 
 if __name__ == "__main__":
-    # Initialize models on startup to reduce first request latency
     try:
+        # Try initializing whisper model on startup
         model_manager.initialize_whisper()
-        model_manager.initialize_llm()
     except Exception as e:
-        logger.warning(f"Initial model loading failed: {str(e)}")
+        logger.warning(f"Initial whisper model loading failed: {str(e)}")
     
     demo = create_demo()
     demo.queue(concurrency_count=1, max_size=5)