Update app.py

app.py

@@ -6,7 +6,8 @@ import tempfile
 import pandas as pd
 import requests
 from bs4 import BeautifulSoup
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer
+from unsloth import FastLanguageModel
 import torch
 import whisper
 from moviepy.editor import VideoFileClip
@@ -15,6 +16,7 @@ import fitz
 import docx
 import yt_dlp
 from functools import lru_cache
+import gc
 
 # Configure logging
 logging.basicConfig(
@@ -44,7 +46,6 @@ class ModelManager:
     def initialize_models(self):
         """Initialize models with optimized settings"""
         try:
-            # Get HuggingFace token
             HUGGINGFACE_TOKEN = os.environ.get('HUGGINGFACE_TOKEN')
             if not HUGGINGFACE_TOKEN:
                 raise ValueError("HUGGINGFACE_TOKEN environment variable not set")
@@ -60,28 +61,45 @@ class ModelManager:
                 use_fast=True,
                 model_max_length=512
             )
-            if self.tokenizer is None:
-                raise RuntimeError("Failed to initialize tokenizer")
             self.tokenizer.pad_token = self.tokenizer.eos_token
 
-            # Load model with optimized memory settings
-            logger.info("Loading model...")
-            self.model = AutoModelForCausalLM.from_pretrained(
-                model_name,
+            # Initialize model with Unsloth optimizations
+            logger.info("Loading model with Unsloth optimizations...")
+            model, tokenizer = FastLanguageModel.from_pretrained(
+                model_name=model_name,
                 token=HUGGINGFACE_TOKEN,
-                torch_dtype=torch.float16,
-                device_map="auto",
-                low_cpu_mem_usage=True,
-                max_memory={0: "6GiB"},
-                load_in_8bit=True
+                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,
+                )
             )
-            if self.model is None:
-                raise RuntimeError("Failed to initialize model")
+            
+            # 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,
+            )
+
+            self.model = model
+            logger.info("Model loaded successfully with Unsloth optimizations")
 
-            # Create pipeline with optimized settings
+            # Create optimized pipeline
             logger.info("Creating pipeline...")
-            self.news_generator = pipeline(
-                "text-generation",
+            self.news_generator = FastLanguageModel.get_pipeline(
                 model=self.model,
                 tokenizer=self.tokenizer,
                 device_map="auto",
@@ -94,18 +112,15 @@ class ModelManager:
                 num_return_sequences=1,
                 early_stopping=True
             )
-            if self.news_generator is None:
-                raise RuntimeError("Failed to initialize news generator pipeline")
 
             # Load Whisper model with optimized settings
             logger.info("Loading Whisper model...")
             self.whisper_model = whisper.load_model(
                 "tiny",
                 device="cuda",
-                download_root="/tmp/whisper"
+                download_root="/tmp/whisper",
+                in_memory=True
             )
-            if self.whisper_model is None:
-                raise RuntimeError("Failed to initialize Whisper model")
 
             logger.info("All models initialized successfully")
             return True
@@ -118,21 +133,31 @@ class ModelManager:
     def reset_models(self):
         """Reset all models and clear GPU memory"""
         try:
-            del self.tokenizer
-            del self.model
-            del self.news_generator
-            del self.whisper_model
+            if hasattr(self, 'model') and self.model is not None:
+                self.model.cpu()
+                del self.model
+            
+            if hasattr(self, 'tokenizer') and self.tokenizer is not None:
+                del self.tokenizer
+                
+            if hasattr(self, 'news_generator') and self.news_generator is not None:
+                del self.news_generator
+                
+            if hasattr(self, 'whisper_model') and self.whisper_model is not None:
+                self.whisper_model.cpu()
+                del self.whisper_model
             
             self.tokenizer = None
             self.model = None
             self.news_generator = None
             self.whisper_model = None
             
-            # Clear CUDA cache
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
                 torch.cuda.synchronize()
             
+            gc.collect()
+            
         except Exception as e:
             logger.error(f"Error during model reset: {str(e)}")
 
@@ -150,12 +175,7 @@ class ModelManager:
 # Create global model manager instance
 model_manager = ModelManager()
 
-# Initialize models at startup
-try:
-    model_manager.initialize_models()
-except Exception as e:
-    logger.error(f"Initial model initialization failed: {str(e)}")
-
+@lru_cache(maxsize=32)
 def download_social_media_video(url):
     """Download a video from social media."""
     ydl_opts = {
@@ -206,7 +226,6 @@ def preprocess_audio(audio_file):
 def transcribe_audio(file):
     """Transcribe an audio or video file."""
     try:
-        # Get initialized models
         _, _, _, whisper_model = model_manager.get_models()
         
         if isinstance(file, str) and file.startswith('http'):
@@ -232,6 +251,7 @@ def transcribe_audio(file):
         logger.error(f"Error transcribing: {str(e)}")
         return f"Error processing the file: {str(e)}"
 
+@lru_cache(maxsize=32)
 def read_document(document_path):
     """Read the content of a document."""
     try:
@@ -251,6 +271,7 @@ def read_document(document_path):
         logger.error(f"Error reading document: {str(e)}")
         return f"Error reading document: {str(e)}"
 
+@lru_cache(maxsize=32)
 def read_url(url):
     """Read the content of a URL."""
     try:
@@ -283,10 +304,8 @@ def process_social_content(url):
 @spaces.GPU(duration=120)
 def generate_news(instructions, facts, size, tone, *args):
     try:
-        # Get initialized models
         tokenizer, _, news_generator, _ = model_manager.get_models()
         
-        # Initialize knowledge base
         knowledge_base = {
             "instructions": instructions,
             "facts": facts,
@@ -296,7 +315,6 @@ def generate_news(instructions, facts, size, tone, *args):
             "social_content": []
         }
 
-        # Parse arguments
         num_audios = 5 * 3
         num_social_urls = 3 * 3
         num_urls = 5
@@ -306,21 +324,18 @@ def generate_news(instructions, facts, size, tone, *args):
         urls = args[num_audios+num_social_urls:num_audios+num_social_urls+num_urls]
         documents = args[num_audios+num_social_urls+num_urls:]
 
-        # Process URLs
         for url in urls:
             if url:
                 content = read_url(url)
                 if content and not content.startswith("Error"):
                     knowledge_base["url_content"].append(content)
 
-        # Process documents
         for document in documents:
             if document is not None:
                 content = read_document(document.name)
                 if content and not content.startswith("Error"):
                     knowledge_base["document_content"].append(content)
 
-        # Process audio files
         for i in range(0, len(audios), 3):
             audio_file, name, position = audios[i:i+3]
             if audio_file is not None:
@@ -330,7 +345,6 @@ def generate_news(instructions, facts, size, tone, *args):
                     "position": position
                 })
 
-        # Process social media content
         for i in range(0, len(social_urls), 3):
             social_url, social_name, social_context = social_urls[i:i+3]
             if social_url:
@@ -344,7 +358,6 @@ def generate_news(instructions, facts, size, tone, *args):
                         "video": social_content["video"]
                     })
 
-        # Build transcriptions
         transcriptions_text = ""
         raw_transcriptions = ""
 
@@ -367,7 +380,7 @@ def generate_news(instructions, facts, size, tone, *args):
         document_content = "\n\n".join(knowledge_base["document_content"])
         url_content = "\n\n".join(knowledge_base["url_content"])
 
-        # Create prompt
+
         prompt = f"""[INST] You are a professional news writer. Write a news article based on the following information:
 
 Instructions: {knowledge_base["instructions"]}
@@ -394,40 +407,45 @@ Follow these requirements:
 
         # Generate article with optimized settings
         with torch.inference_mode():
-            outputs = news_generator(
-                prompt,
-                max_new_tokens=max_tokens,
-                num_return_sequences=1,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.95,
-                repetition_penalty=1.2,
-                early_stopping=True,
-                pad_token_id=tokenizer.eos_token_id
-            )
-        
-        news_article = outputs[0]['generated_text']
-        news_article = news_article.replace('[INST]', '').replace('[/INST]', '').strip()
+            try:
+                news_article = news_generator(
+                    prompt,
+                    max_new_tokens=max_tokens,
+                    num_return_sequences=1,
+                    do_sample=True,
+                    temperature=0.7,
+                    top_p=0.95,
+                    repetition_penalty=1.2,
+                    early_stopping=True
+                )
+                
+                # Process the generated text
+                if isinstance(news_article, list):
+                    news_article = news_article[0]['generated_text']
+                news_article = news_article.replace('[INST]', '').replace('[/INST]', '').strip()
+                
+            except Exception as gen_error:
+                logger.error(f"Error in text generation: {str(gen_error)}")
+                raise
         
         return news_article, raw_transcriptions
 
     except Exception as e:
         logger.error(f"Error generating news: {str(e)}")
         try:
+            # Attempt to recover by resetting and reinitializing models
             model_manager.reset_models()
             model_manager.initialize_models()
-            logger.info("Models reinitialized successfully")
+            logger.info("Models reinitialized successfully after error")
         except Exception as reinit_error:
             logger.error(f"Failed to reinitialize models: {str(reinit_error)}")
         return f"Error generating the news article: {str(e)}", ""
-        
+
 def create_demo():
     with gr.Blocks() as demo:
         gr.Markdown("## Generador de noticias todo en uno")
         
-        # Contenedor principal con dos columnas
         with gr.Row():
-            # Columna izquierda - Formulario principal
             with gr.Column(scale=2):
                 instrucciones = gr.Textbox(
                     label="Instrucciones para la noticia",
@@ -447,14 +465,10 @@ def create_demo():
                     value="neutral"
                 )
 
-            # Columna derecha - Tabs y campos
             with gr.Column(scale=3):
-                # Lista de inputs que empezamos a construir
                 inputs_list = [instrucciones, hechos, tamaño, tono]
 
-                # Tabs en la parte superior
                 with gr.Tabs():
-                    # Audio/Video tabs
                     for i in range(1, 6):
                         with gr.TabItem(f"Audio/Video {i}"):
                             file = gr.File(
@@ -471,7 +485,6 @@ def create_demo():
                             )
                             inputs_list.extend([file, nombre, cargo])
 
-                    # Redes Sociales tabs
                     for i in range(1, 4):
                         with gr.TabItem(f"Red Social {i}"):
                             social_url = gr.Textbox(
@@ -487,7 +500,6 @@ def create_demo():
                             )
                             inputs_list.extend([social_url, social_nombre, social_contexto])
 
-                    # URL tabs
                     for i in range(1, 6):
                         with gr.TabItem(f"URL {i}"):
                             url = gr.Textbox(
@@ -496,7 +508,6 @@ def create_demo():
                             )
                             inputs_list.append(url)
 
-                    # Documento tabs
                     for i in range(1, 6):
                         with gr.TabItem(f"Documento {i}"):
                             documento = gr.File(
@@ -506,10 +517,8 @@ def create_demo():
                             )
                             inputs_list.append(documento)
 
-        # Separador
         gr.Markdown("---")
 
-        # Transcripciones
         with gr.Row():
             transcripciones_output = gr.Textbox(
                 label="Transcripciones",
@@ -517,10 +526,8 @@ def create_demo():
                 show_copy_button=True
             )
 
-        # Separador
         gr.Markdown("---")
 
-        # Botón y output
         with gr.Row():
             generar = gr.Button("Generar borrador")
         
@@ -531,7 +538,6 @@ def create_demo():
                 show_copy_button=True
             )
 
-        # Event handler
         generar.click(
             fn=generate_news,
             inputs=inputs_list,
@@ -540,7 +546,6 @@ def create_demo():
 
     return demo
 
-# Launch the app
 if __name__ == "__main__":
     demo = create_demo()
     demo.queue()