Update app.py

app.py

@@ -17,7 +17,8 @@ from functools import lru_cache
 import gc
 import time
 from huggingface_hub import login
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from unsloth import FastLanguageModel
+from transformers import AutoTokenizer
 
 # Configure logging
 logging.basicConfig(
@@ -51,47 +52,34 @@ class ModelManager:
     
     @spaces.GPU()
     def initialize_llm(self):
-        """Initialize LLM model with standard transformers"""
+        """Initialize LLM model with Unsloth optimization"""
         try:
-            # Use small model for ZeroGPU compatibility
             MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
             
-            logger.info("Loading tokenizer...")
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                MODEL_NAME,
-                token=HUGGINGFACE_TOKEN,
-                use_fast=True
+            logger.info("Loading Unsloth-optimized model...")
+            self.model, self.tokenizer = FastLanguageModel.from_pretrained(
+                model_name = MODEL_NAME,
+                max_seq_length = 2048,
+                dtype = torch.float16,
+                load_in_4bit = True,
+                token = HUGGINGFACE_TOKEN,
             )
             
-            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,
-                # Optimizations for ZeroGPU
-                max_memory={0: "4GB"},
-                offload_folder="offload",
-                offload_state_dict=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=1024
+            # Enable LoRA for better ZeroGPU performance
+            self.model = FastLanguageModel.get_peft_model(
+                self.model,
+                r = 16,
+                target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
+                                 "gate_proj", "up_proj", "down_proj"],
+                lora_alpha = 16,
+                lora_dropout = 0,
+                bias = "none",
+                use_gradient_checkpointing = True,
+                random_state = 3407,
+                max_seq_length = 2048,
             )
             
-            logger.info("LLM initialized successfully")
+            logger.info("LLM initialized successfully with Unsloth")
             self.last_used = time.time()
             return True
             
@@ -101,14 +89,15 @@ class ModelManager:
 
     @spaces.GPU()
     def initialize_whisper(self):
-        """Initialize Whisper model for audio transcription"""
+        """Initialize Whisper model with safety fix"""
         try:
             logger.info("Loading Whisper model...")
-            # Using tiny model for efficiency but can be changed based on needs
+            # Load with weights_only=True for security
             self.whisper_model = whisper.load_model(
                 "tiny",
                 device="cuda" if torch.cuda.is_available() else "cpu",
-                download_root="/tmp/whisper"
+                download_root="/tmp/whisper",
+                weights_only=True  # Security fix
             )
             logger.info("Whisper model initialized successfully")
             self.last_used = time.time()
@@ -119,7 +108,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 or self.pipeline is None:
+        if self.tokenizer is None or self.model is None:
             logger.info("LLM not initialized, initializing...")
             self.initialize_llm()
         self.last_used = time.time()
@@ -134,26 +123,21 @@ class ModelManager:
     def reset_models(self, force=False):
         """Reset models to free memory if they haven't been used recently"""
         current_time = time.time()
-        # Only reset if forced or models haven't been used for 10 minutes
         if force or (current_time - self.last_used > 600):  
             try:
                 logger.info("Resetting models to free memory...")
                 
-                if hasattr(self, 'model') and self.model is not None:
+                if self.model is not None:
                     del self.model
                     
-                if hasattr(self, 'tokenizer') and self.tokenizer is not None:
+                if 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:
+                if 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():
@@ -166,7 +150,6 @@ class ModelManager:
             except Exception as e:
                 logger.error(f"Error resetting models: {str(e)}")
 
-# Create global model manager instance
 model_manager = ModelManager()
 
 @lru_cache(maxsize=32)
@@ -197,7 +180,6 @@ def convert_video_to_audio(video_file):
         with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_file:
             output_file = temp_file.name
         
-        # Use ffmpeg directly via subprocess
         command = [
             "ffmpeg", 
             "-i", video_file, 
@@ -205,7 +187,7 @@ def convert_video_to_audio(video_file):
             "-map", "a",
             "-vn",
             output_file,
-            "-y"  # Overwrite output file if it exists
+            "-y"
         ]
         
         subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
@@ -239,10 +221,10 @@ def transcribe_audio(file):
             file_path = download_social_media_video(file)
         elif isinstance(file, str) and file.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
             file_path = convert_video_to_audio(file)
-        elif file is not None:  # Handle file object from Gradio
+        elif file is not None:
             file_path = preprocess_audio(file.name)
         else:
-            return ""  # Return empty string for None input
+            return ""
 
         logger.info(f"Transcribing audio: {file_path}")
         if not os.path.exists(file_path):
@@ -250,13 +232,10 @@ def transcribe_audio(file):
 
         with torch.inference_mode():
             result = model_manager.whisper_model.transcribe(file_path)
-            if not result:
-                raise RuntimeError("Transcription failed to produce results")
                 
         transcription = result.get("text", "Error in transcription")
         logger.info(f"Transcription completed: {transcription[:50]}...")
         
-        # Clean up temp file
         try:
             if os.path.exists(file_path):
                 os.remove(file_path)
@@ -302,22 +281,19 @@ def read_url(url):
         response.raise_for_status()
         soup = BeautifulSoup(response.content, 'html.parser')
         
-        # Remove non-content elements
         for element in soup(["script", "style", "meta", "noscript", "iframe", "header", "footer", "nav"]):
             element.extract()
             
-        # Extract main content
         main_content = soup.find("main") or soup.find("article") or soup.find("div", class_=["content", "main", "article"])
         if main_content:
             text = main_content.get_text(separator='\n', strip=True)
         else:
             text = soup.get_text(separator='\n', strip=True)
             
-        # Clean up whitespace
         lines = [line.strip() for line in text.split('\n') if line.strip()]
         text = '\n'.join(lines)
         
-        return text[:10000]  # Limit to 10k chars to avoid huge inputs
+        return text[:10000]
     except Exception as e:
         logger.error(f"Error reading URL: {str(e)}")
         return f"Error reading URL: {str(e)}"
@@ -347,16 +323,13 @@ def process_social_content(url):
 def generate_news(instructions, facts, size, tone, *args):
     """Generate a news article based on provided data"""
     try:
-        # Ensure size is integer
         if isinstance(size, float):
             size = int(size)
         elif not isinstance(size, int):
-            size = 250  # Default size
+            size = 250
             
-        # Check if models are initialized
         model_manager.check_llm_initialized()
         
-        # Prepare data structure for inputs
         knowledge_base = {
             "instructions": instructions or "",
             "facts": facts or "",
@@ -366,15 +339,12 @@ def generate_news(instructions, facts, size, tone, *args):
             "social_content": []
         }
 
-        # Define the indices for parsing args
         num_audios = 5 * 3
         num_social_urls = 3 * 3
         num_urls = 5
 
-        # Parse arguments
-        args = list(args)  # Convert tuple to list for easier manipulation
+        args = list(args)
         
-        # Ensure we have enough arguments
         while len(args) < (num_audios + num_social_urls + num_urls + 5):
             args.append("")
             
@@ -383,7 +353,6 @@ 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 with progress reporting
         logger.info("Processing URLs...")
         for url in urls:
             if url and isinstance(url, str) and url.strip():
@@ -391,7 +360,6 @@ def generate_news(instructions, facts, size, tone, *args):
                 if content and not content.startswith("Error"):
                     knowledge_base["url_content"].append(content)
 
-        # Process documents
         logger.info("Processing documents...")
         for document in documents:
             if document and hasattr(document, 'name'):
@@ -399,10 +367,9 @@ def generate_news(instructions, facts, size, tone, *args):
                 if content and not content.startswith("Error"):
                     knowledge_base["document_content"].append(content)
 
-        # Process audio/video files
         logger.info("Processing audio/video files...")
         for i in range(0, len(audios), 3):
-            if i+2 < len(audios):  # Ensure we have complete set of 3 elements
+            if i+2 < len(audios):
                 audio_file, name, position = audios[i:i+3]
                 if audio_file and hasattr(audio_file, 'name'):
                     knowledge_base["audio_data"].append({
@@ -411,10 +378,9 @@ def generate_news(instructions, facts, size, tone, *args):
                         "position": position or "Not specified"
                     })
 
-        # Process social media content
         logger.info("Processing social media content...")
         for i in range(0, len(social_urls), 3):
-            if i+2 < len(social_urls):  # Ensure we have complete set of 3 elements
+            if i+2 < len(social_urls):
                 social_url, social_name, social_context = social_urls[i:i+3]
                 if social_url and isinstance(social_url, str) and social_url.strip():
                     social_content = process_social_content(social_url)
@@ -427,11 +393,9 @@ def generate_news(instructions, facts, size, tone, *args):
                             "video": social_content.get("video", "")
                         })
 
-        # Prepare transcriptions text
         transcriptions_text = ""
         raw_transcriptions = ""
 
-        # Process audio data transcriptions
         logger.info("Transcribing audio...")
         for idx, data in enumerate(knowledge_base["audio_data"]):
             if data["audio"] is not None:
@@ -440,10 +404,8 @@ def generate_news(instructions, facts, size, tone, *args):
                     transcriptions_text += f'"{transcription}" - {data["name"]}, {data["position"]}\n\n'
                     raw_transcriptions += f'[Audio/Video {idx + 1}]: "{transcription}" - {data["name"]}, {data["position"]}\n\n'
 
-        # Process social media content transcriptions
         for idx, data in enumerate(knowledge_base["social_content"]):
             if data["text"] and not str(data["text"]).startswith("Error"):
-                # Truncate long texts for the prompt
                 text_excerpt = data["text"][:500] + "..." if len(data["text"]) > 500 else data["text"]
                 social_text = f'[Social media {idx+1} - text]: "{text_excerpt}" - {data["name"]}, {data["context"]}\n\n'
                 transcriptions_text += social_text
@@ -454,10 +416,8 @@ def generate_news(instructions, facts, size, tone, *args):
                 transcriptions_text += video_transcription
                 raw_transcriptions += video_transcription
 
-        # Combine document content and URL content (with truncation for very long content)
         document_summaries = []
         for idx, doc in enumerate(knowledge_base["document_content"]):
-            # Truncate long documents
             if len(doc) > 1000:
                 doc_excerpt = doc[:1000] + "... [document continues]"
             else:
@@ -468,7 +428,6 @@ def generate_news(instructions, facts, size, tone, *args):
         
         url_summaries = []
         for idx, url_content in enumerate(knowledge_base["url_content"]):
-            # Truncate long URL content
             if len(url_content) > 1000:
                 url_excerpt = url_content[:1000] + "... [content continues]"
             else:
@@ -477,7 +436,6 @@ def generate_news(instructions, facts, size, tone, *args):
             
         url_content = "\n\n".join(url_summaries)
 
-        # Create prompt for the model
         prompt = f"""<s>[INST] You are a professional news writer. Write a news article based on the following information:
 
 Instructions: {knowledge_base["instructions"]}
@@ -504,35 +462,33 @@ Follow these requirements:
 - Do not invent information
 - Be rigorous with the provided facts [/INST]"""
 
-        # 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(
+            inputs = model_manager.tokenizer(
                 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
+                return_tensors = "pt",
+                padding = True,
+                truncation = True,
+                max_length = 2048,
+            ).to("cuda")
+
+            outputs = model_manager.model.generate(
+                **inputs,
+                max_new_tokens = size + 100,
+                temperature = 0.7,
+                do_sample = True,
+                pad_token_id = model_manager.tokenizer.eos_token_id,
             )
             
-            # Extract generated text
-            generated_text = outputs[0]['generated_text']
+            generated_text = model_manager.tokenizer.decode(outputs[0], skip_special_tokens = True)
             
-            # 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)
+                prompt_fragment = " ".join(prompt.split()[:50])
                 if prompt_fragment in generated_text:
                     news_article = generated_text[generated_text.find(prompt_fragment) + len(prompt_fragment):].strip()
                 else:
@@ -549,173 +505,149 @@ Follow these requirements:
     except Exception as e:
         logger.error(f"Error generating news: {str(e)}")
         try:
-            # Reset models to recover from errors
             model_manager.reset_models(force=True)
         except Exception as reset_error:
             logger.error(f"Failed to reset models: {str(reset_error)}")
-        return f"Error generando la noticia: {str(e)}", "Error procesando las transcripciones."
+        return f"Error generating news: {str(e)}", "Error processing transcriptions."
 
 def create_demo():
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 📰 NewsIA - Generador de Noticias IA")
-        gr.Markdown("Crea noticias profesionales a partir de múltiples fuentes de información.")
+        gr.Markdown("# 📰 NewsIA - AI News Generator")
+        gr.Markdown("Create professional news articles from multiple sources.")
         
         with gr.Row():
             with gr.Column(scale=2):
-                instrucciones = gr.Textbox(
-                    label="Instrucciones para la noticia",
-                    placeholder="Escribe instrucciones específicas para la generación de tu noticia",
-                    lines=2,
-                    value=""
+                instructions = gr.Textbox(
+                    label="News Instructions",
+                    placeholder="Enter specific instructions for news generation",
+                    lines=2
                 )
-                hechos = gr.Textbox(
-                    label="Hechos principales",
-                    placeholder="Describe los hechos más importantes que debe incluir la noticia",
-                    lines=4,
-                    value=""
+                facts = gr.Textbox(
+                    label="Key Facts",
+                    placeholder="Describe the most important facts to include",
+                    lines=4
                 )
                 
                 with gr.Row():
-                    tamaño = gr.Slider(
-                        label="Longitud aproximada (palabras)",
+                    size = gr.Slider(
+                        label="Approximate Length (words)",
                         minimum=100,
                         maximum=500,
                         value=250,
                         step=50
                     )
-                    tono = gr.Dropdown(
-                        label="Tono de la noticia",
-                        choices=["serio", "neutral", "divertido", "formal", "informal", "urgente"],
+                    tone = gr.Dropdown(
+                        label="News Tone",
+                        choices=["serious", "neutral", "funny", "formal", "informal", "urgent"],
                         value="neutral"
                     )
 
             with gr.Column(scale=3):
-                # Inicializamos la lista de inputs con valores conocidos
                 inputs_list = []
-                inputs_list.append(instrucciones)
-                inputs_list.append(hechos)
-                inputs_list.append(tamaño)
-                inputs_list.append(tono)
+                inputs_list.extend([instructions, facts, size, tone])
 
                 with gr.Tabs():
-                    with gr.TabItem("📝 Documentos"):
-                        documentos = []
-                        for i in range(1, 6):  # Mantenemos 5 documentos como en el original
-                            documento = gr.File(
-                                label=f"Documento {i}",
+                    with gr.TabItem("📝 Documents"):
+                        documents = []
+                        for i in range(1, 6):
+                            doc = gr.File(
+                                label=f"Document {i}",
                                 file_types=["pdf", "docx", "xlsx", "csv"],
-                                file_count="single",
-                                value=None
+                                file_count="single"
                             )
-                            documentos.append(documento)
-                            inputs_list.append(documento)
+                            documents.append(doc)
+                            inputs_list.append(doc)
 
                     with gr.TabItem("🔊 Audio/Video"):
-                        for i in range(1, 6):  # Mantenemos 5 fuentes como en el original
+                        for i in range(1, 6):
                             with gr.Group():
-                                gr.Markdown(f"**Fuente {i}**")
+                                gr.Markdown(f"**Source {i}**")
                                 file = gr.File(
                                     label=f"Audio/Video {i}",
-                                    file_types=["audio", "video"],
-                                    value=None
+                                    file_types=["audio", "video"]
                                 )
                                 with gr.Row():
-                                    nombre = gr.Textbox(
-                                        label="Nombre",
-                                        placeholder="Nombre del entrevistado",
-                                        value=""
+                                    name = gr.Textbox(
+                                        label="Name",
+                                        placeholder="Interviewee name"
                                     )
-                                    cargo = gr.Textbox(
-                                        label="Cargo/Rol",
-                                        placeholder="Cargo o rol",
-                                        value=""
+                                    position = gr.Textbox(
+                                        label="Position/Role",
+                                        placeholder="Position or role"
                                     )
-                                inputs_list.append(file)
-                                inputs_list.append(nombre)
-                                inputs_list.append(cargo)
+                                inputs_list.extend([file, name, position])
 
                     with gr.TabItem("🌐 URLs"):
-                        for i in range(1, 6):  # Mantenemos 5 URLs como en el original
+                        for i in range(1, 6):
                             url = gr.Textbox(
                                 label=f"URL {i}",
-                                placeholder="https://...",
-                                value=""
+                                placeholder="https://..."
                             )
                             inputs_list.append(url)
 
-                    with gr.TabItem("📱 Redes Sociales"):
-                        for i in range(1, 4):  # Mantenemos 3 redes sociales como en el original
+                    with gr.TabItem("📱 Social Media"):
+                        for i in range(1, 4):
                             with gr.Group():
-                                gr.Markdown(f"**Red Social {i}**")
+                                gr.Markdown(f"**Social Media {i}**")
                                 social_url = gr.Textbox(
-                                    label=f"URL",
-                                    placeholder="https://...",
-                                    value=""
+                                    label="URL",
+                                    placeholder="https://..."
                                 )
                                 with gr.Row():
-                                    social_nombre = gr.Textbox(
-                                        label=f"Nombre/Cuenta",
-                                        placeholder="Nombre de la persona o cuenta",
-                                        value=""
+                                    social_name = gr.Textbox(
+                                        label="Account/Name",
+                                        placeholder="Account or person name"
                                     )
-                                    social_contexto = gr.Textbox(
-                                        label=f"Contexto",
-                                        placeholder="Contexto relevante",
-                                        value=""
+                                    social_context = gr.Textbox(
+                                        label="Context",
+                                        placeholder="Relevant context"
                                     )
-                                inputs_list.append(social_url)
-                                inputs_list.append(social_nombre)
-                                inputs_list.append(social_contexto)
+                                inputs_list.extend([social_url, social_name, social_context])
 
         with gr.Row():
-            generar = gr.Button("✨ Generar Noticia", variant="primary")
-            reset = gr.Button("🔄 Limpiar Todo")
+            generate_btn = gr.Button("✨ Generate News", variant="primary")
+            reset_btn = gr.Button("🔄 Clear All")
 
         with gr.Tabs():
-            with gr.TabItem("📄 Noticia Generada"):
-                noticia_output = gr.Textbox(
-                    label="Borrador de la noticia",
+            with gr.TabItem("📄 Generated News"):
+                news_output = gr.Textbox(
+                    label="News Draft",
                     lines=15,
-                    show_copy_button=True,
-                    value=""
+                    show_copy_button=True
                 )
             
-            with gr.TabItem("🎙️ Transcripciones"):
-                transcripciones_output = gr.Textbox(
-                    label="Transcripciones de fuentes",
+            with gr.TabItem("🎙️ Transcriptions"):
+                transcriptions_output = gr.Textbox(
+                    label="Source Transcriptions",
                     lines=10,
-                    show_copy_button=True,
-                    value=""
+                    show_copy_button=True
                 )
 
-        # Set up event handlers
-        generar.click(
+        generate_btn.click(
             fn=generate_news,
             inputs=inputs_list,
-            outputs=[noticia_output, transcripciones_output]
+            outputs=[news_output, transcriptions_output]
         )
         
-        # Reset functionality to clear all inputs
         def reset_all():
-            return [""] * len(inputs_list) + ["", ""]
+            return [None]*len(inputs_list) + ["", ""]
         
-        reset.click(
+        reset_btn.click(
             fn=reset_all,
             inputs=None,
-            outputs=inputs_list + [noticia_output, transcripciones_output]
+            outputs=inputs_list + [news_output, transcriptions_output]
         )
 
     return demo
 
 if __name__ == "__main__":
     try:
-        # Try initializing whisper model on startup
         model_manager.initialize_whisper()
     except Exception as e:
         logger.warning(f"Initial whisper model loading failed: {str(e)}")
     
     demo = create_demo()
-    demo.queue(concurrency_count=1, max_size=5)
+    demo.queue(max_size=5)
     demo.launch(
         share=True,
         server_name="0.0.0.0",