Update app.py

app.py

@@ -3,267 +3,117 @@ from transformers import AutoTokenizer
 from peft import AutoPeftModelForSeq2SeqLM
 import torch
 
-# Define the Hugging Face repository ID
+# Model configuration
 repo_id = "sairika/FLAN-T5-Base-dialogsum-lora"
 
-# Sample dialogues for examples
-SAMPLE_DIALOGUES = [
-    """Alice: Hi Bob, how was your meeting with the client today?
-Bob: It went really well! They loved our proposal and want to move forward.
-Alice: That's fantastic news! What's the next step?
-Bob: We need to prepare the contract and schedule a follow-up meeting next week.
-Alice: Great, I'll help you with the contract preparation.""",
+# Sample dialogues
+examples = [
+    "Alice: How was your meeting today?\nBob: It went great! The client approved our proposal.\nAlice: That's wonderful news!\nBob: Yes, we start next Monday.",
     
-    """Customer: I'm having trouble with my internet connection. It keeps dropping out.
-Support: I'm sorry to hear that. Let me help you troubleshoot this issue.
-Customer: I've already tried restarting my router.
-Support: Okay, let's check your signal strength. Can you tell me what lights are showing on your modem?
-Customer: There are green lights for power and internet, but the wifi light is blinking orange.
-Support: That indicates a wifi connectivity issue. Let's reset your wifi settings.""",
+    "Customer: My internet isn't working.\nSupport: I can help with that. Have you tried restarting your router?\nCustomer: Yes, but it's still not working.\nSupport: Let me check your connection status.",
     
-    """Mom: Have you finished your homework yet?
-Child: Almost done! I just have math left.
-Mom: Do you need any help with it?
-Child: Actually yes, I'm stuck on these algebra problems.
-Mom: Let me take a look. Oh, these are quadratic equations. Remember the formula we practiced?
-Child: Oh right! ax² + bx + c = 0. Thanks mom!"""
+    "Teacher: Did you finish your homework?\nStudent: Almost done, just the math problems left.\nTeacher: Need help with anything?\nStudent: Yes, I'm stuck on question 5."
 ]
 
-# Load the tokenizer and the PEFT model
-try:
-    tokenizer = AutoTokenizer.from_pretrained(repo_id)
-    model = AutoPeftModelForSeq2SeqLM.from_pretrained(
-        repo_id, 
-        device_map="auto", 
-        torch_dtype=torch.bfloat16
-    )
-    print(f"✅ Model and tokenizer loaded successfully from {repo_id}")
-    
-    def create_prompts(dialogues, model_type):
-        """Create appropriate prompts based on model type"""
-        if model_type in ['flan-t5', 't5']:
-            prompts = [f"Summarize the following conversation.\n\n{dialogue}\n\nSummary: "
-                      for dialogue in dialogues]
-        else:
-            prompts = dialogues
-        return prompts
-    
-    def generate_summary(dialogue, max_length, num_beams, length_penalty):
-        """Generates a summary for a given dialogue with customizable parameters"""
-        if not dialogue.strip():
-            return "⚠️ Please enter a dialogue to summarize."
-        
-        try:
-            model_type = 'flan-t5' if 'flan-t5' in repo_id else 't5' if 't5' in repo_id else 'bart'
-            prompt = create_prompts([dialogue], model_type)[0]
-            
-            inputs = tokenizer(
-                prompt,
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=512
-            ).to(model.device)
-            
-            model.eval()
-            with torch.no_grad():
-                outputs = model.generate(
-                    **inputs,
-                    max_new_tokens=int(max_length),
-                    num_beams=int(num_beams),
-                    length_penalty=float(length_penalty),
-                    early_stopping=True,
-                    do_sample=False
-                )
-            
-            summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            return f"📝 **Summary:** {summary}"
-            
-        except Exception as e:
-            return f"❌ Error generating summary: {str(e)}"
+def load_model():
+    """Load the model and tokenizer"""
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(repo_id)
+        model = AutoPeftModelForSeq2SeqLM.from_pretrained(
+            repo_id, 
+            device_map="auto", 
+            torch_dtype=torch.bfloat16
+        )
+        return tokenizer, model, True
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None, None, False
+
+def summarize_dialogue(dialogue):
+    """Generate summary for the dialogue"""
+    if not dialogue.strip():
+        return "Please enter a dialogue to summarize."
     
-    # Custom CSS for better styling
-    custom_css = """
-    .gradio-container {
-        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-    }
-    .header {
-        text-align: center;
-        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
-        color: white;
-        padding: 2rem;
-        border-radius: 10px;
-        margin-bottom: 2rem;
-    }
-    .footer {
-        text-align: center;
-        margin-top: 2rem;
-        padding: 1rem;
-        background-color: #f8f9fa;
-        border-radius: 10px;
-    }
-    """
+    if not model_loaded:
+        return "Error: Model not loaded. Please check the logs."
     
-    # Create the Gradio interface
-    with gr.Blocks(css=custom_css, title="Dialogue Summarization Demo") as iface:
-        # Header
-        gr.HTML("""
-        <div class="header">
-            <h1>🤖 AI Dialogue Summarization</h1>
-            <p>Transform lengthy conversations into concise, meaningful summaries using a fine-tuned FLAN-T5 model</p>
-        </div>
-        """)
-        
-        with gr.Row():
-            with gr.Column(scale=2):
-                dialogue_input = gr.Textbox(
-                    label="📝 Enter Your Dialogue",
-                    placeholder="Paste your conversation here...\n\nExample:\nPerson A: Hello, how are you?\nPerson B: I'm doing well, thanks for asking!",
-                    lines=10,
-                    max_lines=20
-                )
-                
-                with gr.Row():
-                    summarize_btn = gr.Button("✨ Generate Summary", variant="primary", size="lg")
-                    clear_btn = gr.Button("🗑️ Clear", variant="secondary")
-            
-            with gr.Column(scale=1):
-                gr.Markdown("### ⚙️ Generation Parameters")
-                max_length = gr.Slider(
-                    minimum=50,
-                    maximum=256,
-                    value=128,
-                    step=10,
-                    label="Max Summary Length",
-                    info="Maximum number of tokens in the summary"
-                )
-                num_beams = gr.Slider(
-                    minimum=1,
-                    maximum=8,
-                    value=4,
-                    step=1,
-                    label="Number of Beams",
-                    info="Higher values = better quality but slower"
-                )
-                length_penalty = gr.Slider(
-                    minimum=0.1,
-                    maximum=2.0,
-                    value=0.6,
-                    step=0.1,
-                    label="Length Penalty",
-                    info="Controls summary length preference"
-                )
-        
-        # Output
-        summary_output = gr.Textbox(
-            label="🎯 Generated Summary",
-            lines=5,
-            show_copy_button=True
-        )
-        
-        # Examples section
-        gr.Markdown("### 💡 Try These Examples")
-        gr.Examples(
-            examples=SAMPLE_DIALOGUES,
-            inputs=dialogue_input,
-            label="Click on any example to load it:"
-        )
+    try:
+        # Create prompt
+        prompt = f"Summarize the following conversation.\n\n{dialogue}\n\nSummary: "
         
-        # Model information
-        gr.Markdown(f"""
-        ### 📊 Model Information
-        - **Model:** {repo_id}
-        - **Base Architecture:** FLAN-T5-Base with LoRA fine-tuning
-        - **Task:** Dialogue Summarization
-        - **Training Dataset:** DialogSum
-        """)
+        # Tokenize input
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=512
+        ).to(model.device)
         
-        # Footer
-        gr.HTML("""
-        <div class="footer">
-            <p>Built with ❤️ using Gradio and Hugging Face Transformers</p>
-            <p><em>This demo showcases AI-powered dialogue summarization capabilities</em></p>
-        </div>
-        """)
+        # Generate summary
+        model.eval()
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=128,
+                num_beams=4,
+                length_penalty=0.6,
+                early_stopping=True,
+                do_sample=False
+            )
         
-        # Event handlers
-        summarize_btn.click(
-            fn=generate_summary,
-            inputs=[dialogue_input, max_length, num_beams, length_penalty],
-            outputs=summary_output
-        )
+        # Decode and return
+        summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return summary
         
-        clear_btn.click(
-            fn=lambda: ("", ""),
-            outputs=[dialogue_input, summary_output]
-        )
-    
-    # Launch the interface
-    iface.launch(
-        share=True,  # Creates a public link
-        show_error=True,
-        show_api=False,
-        favicon_path=None,
-        ssl_verify=False
-    )
+    except Exception as e:
+        return f"Error generating summary: {str(e)}"
 
-except Exception as e:
-    print(f"❌ Error loading model or setting up Gradio: {e}")
+# Load model
+tokenizer, model, model_loaded = load_model()
+
+# Create Gradio interface
+with gr.Blocks(title="Dialogue Summarization") as demo:
+    gr.Markdown("# 🤖 Dialogue Summarization")
+    gr.Markdown("Enter a conversation below to generate an AI summary using FLAN-T5.")
     
-    # Enhanced error interface
-    def error_message(dialogue):
-        return f"""
-        ❌ **Error Loading Model**
-        
-        Unfortunately, there was an error loading the model from `{repo_id}`.
-        
-        **Possible causes:**
-        - Model repository not accessible
-        - Insufficient memory/resources  
-        - Network connectivity issues
-        - Invalid model format
-        
-        **Error details:** {str(e)}
+    with gr.Row():
+        with gr.Column():
+            dialogue_input = gr.Textbox(
+                label="Dialogue",
+                placeholder="Enter your conversation here...",
+                lines=8
+            )
+            
+            submit_btn = gr.Button("Generate Summary", variant="primary")
+            clear_btn = gr.Button("Clear")
         
-        Please check the logs and try again later.
-        """
+        with gr.Column():
+            summary_output = gr.Textbox(
+                label="Summary",
+                lines=6,
+                show_copy_button=True
+            )
     
-    custom_css_error = """
-    .error-container {
-        background-color: #fee;
-        border: 2px solid #fcc;
-        border-radius: 10px;
-        padding: 2rem;
-        margin: 2rem 0;
-    }
-    """
+    # Examples
+    gr.Examples(
+        examples=examples,
+        inputs=dialogue_input,
+        label="Try these examples:"
+    )
     
-    with gr.Blocks(css=custom_css_error, title="Model Loading Error") as error_iface:
-        gr.HTML("""
-        <div class="error-container">
-            <h1>🚨 Model Loading Error</h1>
-            <p>There was an issue loading the dialogue summarization model.</p>
-        </div>
-        """)
-        
-        dialogue_input = gr.Textbox(
-            label="Enter Dialogue",
-            placeholder="The model could not be loaded...",
-            interactive=False
-        )
-        
-        error_output = gr.Textbox(
-            label="Error Details",
-            value=error_message(""),
-            interactive=False
-        )
-        
-        gr.Markdown(f"""
-        ### 🔧 Troubleshooting Steps:
-        1. Check if the model repository `{repo_id}` exists and is accessible
-        2. Verify you have sufficient system resources (RAM/GPU)
-        3. Ensure all required dependencies are installed
-        4. Check network connectivity to Hugging Face Hub
-        """)
+    # Event handlers
+    submit_btn.click(
+        fn=summarize_dialogue,
+        inputs=dialogue_input,
+        outputs=summary_output
+    )
     
-    error_iface.launch(share=True)
+    clear_btn.click(
+        fn=lambda: ("", ""),
+        outputs=[dialogue_input, summary_output]
+    )
+
+# Launch
+if __name__ == "__main__":
+    demo.launch(share=True)