Update app.py

app.py

@@ -1,7 +1,4 @@
 import gradio as gr
-from pydub import AudioSegment
-from google import genai
-from google.genai import types
 import random
 import time
 import os
@@ -14,11 +11,13 @@ import edge_tts
 import asyncio
 import aiofiles
 import mimetypes
-from typing import List, Dict
+from typing import List
 
+# New import for PDF parsing
+from PyPDF2 import PdfReader
 
 # Define model name clearly
-MODEL_NAME = "unsloth/gemma-3-1b-pt" #HuggingFaceH4/zephyr-7b-alpha"
+MODEL_NAME = "unsloth/gemma-3-1b-pt"  # HuggingFaceH4/zephyr-7b-alpha
 
 # Device setup
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -27,7 +26,7 @@ print(f"Using device: {device}")
 # Load model and tokenizer (explicit evaluation mode)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME, 
+    MODEL_NAME,
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
 ).eval().to(device)
 
@@ -97,7 +96,7 @@ You are a professional podcast generator. Your task is to generate a professiona
 Follow this example structure:
 {example}
 """
-        user_prompt = ""
+        # Build the user prompt
         if prompt and file_obj:
             user_prompt = f"Please generate a podcast script based on the uploaded file following user input:\n{prompt}"
         elif prompt:
@@ -105,42 +104,35 @@ Follow this example structure:
         else:
             user_prompt = "Please generate a podcast script based on the uploaded file."
 
-        messages = []
-        
-        # If file is provided, add it to the messages
+        # If a file is provided, extract its text and append
         if file_obj:
-            file_data = await self._read_file_bytes(file_obj)
-            mime_type = self._get_mime_type(file_obj.name)
-            
-            messages.append(
-                types.Content(
-                    role="user",
-                    parts=[
-                        types.Part.from_bytes(
-                            data=file_data,
-                            mime_type=mime_type,
-                        )
-                    ],
-                )
-            )
-        
-        # Add text prompt
-        messages.append(
-            types.Content(
-                role="user",
-                parts=[
-                    types.Part.from_text(text=user_prompt)
-                ],
-            )
-        )
+            # enforce size limit
+            file_size = getattr(file_obj, 'size', os.path.getsize(file_obj.name))
+            if file_size > MAX_FILE_SIZE_BYTES:
+                raise Exception(f"File size exceeds the {MAX_FILE_SIZE_MB}MB limit. Please upload a smaller file.")
+
+            # extract text based on mime
+            ext = os.path.splitext(file_obj.name)[1].lower()
+            if ext == '.pdf':
+                reader = PdfReader(file_obj)
+                text = "\n\n".join(page.extract_text() or '' for page in reader.pages)
+            else:
+                # txt or other
+                if hasattr(file_obj, 'read'):
+                    raw = file_obj.read()
+                else:
+                    raw = await aiofiles.open(file_obj.name, 'rb').read()
+                text = raw.decode(errors='ignore')
+
+            user_prompt += f"\n\n―― FILE CONTENT ――\n{text}"
+
+        # Combine system and user prompts
+        prompt_text = system_prompt + "\n" + user_prompt
 
         try:
             if progress:
                 progress(0.3, "Generating podcast script...")
 
-            # Compose the prompt from your messages
-            prompt_text = system_prompt + "\n" + "\n".join([msg["content"] for msg in messages])
-
             def hf_generate(prompt_text):
                 inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
                 outputs = model.generate(
@@ -149,8 +141,7 @@ Follow this example structure:
                     do_sample=True,
                     temperature=1.0
                 )
-                text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-                return text
+                return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
             generated_text = await asyncio.wait_for(
                 asyncio.to_thread(hf_generate, prompt_text),
@@ -162,14 +153,14 @@ Follow this example structure:
         except Exception as e:
             raise Exception(f"Failed to generate podcast script: {e}")
 
-        print(f"Generated podcast script:\n{generated_text}")
-
         if progress:
             progress(0.4, "Script generated successfully!")
 
-        # Ensure the return type matches the original code (as JSON)
         return json.loads(generated_text)
 
+    # ... rest of class unchanged ...
+
+
     
     async def _read_file_bytes(self, file_obj) -> bytes:
         """Read file bytes from a file object"""