Update app.py

app.py

@@ -1,4 +1,3 @@
-# Remove the pip install line - use requirements.txt instead
 import torch
 from transformers import (
     AutoTokenizer,
@@ -15,6 +14,7 @@ from typing import List, Tuple
 
 # Create cache directory
 os.makedirs("model_cache", exist_ok=True)
+os.makedirs("examples", exist_ok=True)  # Create examples directory
 
 # Configuration for 4-bit quantization
 quant_config = BitsAndBytesConfig(
@@ -54,12 +54,6 @@ class RiverPollutionAnalyzer:
         except Exception as e:
             raise RuntimeError(f"Model loading failed: {str(e)}")
 
-        # [Rest of your class implementation remains unchanged]
-        # ... (keep all your existing methods) ...
-
-# [Rest of your Gradio implementation remains unchanged]
-# ... (keep all your existing UI code) ...
-
         self.pollutants = [
             "plastic waste", "chemical foam", "industrial discharge",
             "sewage water", "oil spill", "organic debris",
@@ -110,22 +104,97 @@ Severity: [number]"""
         except Exception as e:
             return f"⚠️ Analysis failed: {str(e)}"
 
-    # [Keep all your existing parsing/formatting methods unchanged]
     def _parse_response(self, analysis: str) -> Tuple[List[str], int]:
-        """Same parsing logic as before"""
-        # ... (unchanged from your original code) ...
+        """Parse the model response into pollutants list and severity score"""
+        pollutants = []
+        severity = 0
+        
+        # Extract pollutants
+        pollutants_match = re.search(r"Pollutants:\s*\[(.*?)\]", analysis)
+        if pollutants_match:
+            pollutants_str = pollutants_match.group(1)
+            pollutants = [p.strip() for p in pollutants_str.split(",") if p.strip()]
+        
+        # Extract severity
+        severity_match = re.search(r"Severity:\s*(\d+)", analysis)
+        if severity_match:
+            severity = int(severity_match.group(1))
+        
+        # If parsing failed, fallback to calculating severity
+        if not severity or severity < 1 or severity > 10:
+            severity = self._calculate_severity(pollutants)
+            
+        return pollutants, severity
     
     def _calculate_severity(self, pollutants: List[str]) -> int:
-        """Same severity calculation"""
-        # ... (unchanged from your original code) ...
+        """Calculate severity based on pollutants"""
+        if not pollutants:
+            return 1
+        
+        severity_map = {
+            "plastic waste": 4,
+            "chemical foam": 7,
+            "industrial discharge": 8,
+            "sewage water": 6,
+            "oil spill": 9,
+            "organic debris": 3,
+            "construction waste": 5,
+            "medical waste": 8,
+            "floating trash": 4,
+            "algal bloom": 6,
+            "toxic sludge": 9,
+            "agricultural runoff": 5
+        }
+        
+        base_score = sum(severity_map.get(p, 3) for p in pollutants)
+        avg_score = base_score / len(pollutants)
+        return min(10, max(1, round(avg_score)))
     
     def _format_analysis(self, pollutants: List[str], severity: int) -> str:
-        """Same formatting"""
-        # ... (unchanged from your original code) ...
+        """Format the analysis results into a markdown report"""
+        if not pollutants:
+            pollutants = ["No visible pollution detected"]
+            
+        pollutants_list = "\n".join(f"- {p}" for p in pollutants)
+        severity_desc = self.severity_descriptions.get(severity, "Unknown severity level")
+        
+        return f"""
+## Pollution Analysis Report
+
+### Identified Pollutants:
+{pollutants_list}
+
+### Severity Assessment:
+**Level {severity}/10** - {severity_desc}
+
+### Recommended Actions:
+{self._get_recommendations(severity)}
+"""
+    
+    def _get_recommendations(self, severity: int) -> str:
+        """Get recommendations based on severity level"""
+        if severity <= 3:
+            return "Monitor the situation. Consider community clean-up efforts."
+        elif severity <= 5:
+            return "Local authorities should investigate. Basic remediation needed."
+        elif severity <= 7:
+            return "Immediate containment required. Environmental assessment needed."
+        elif severity <= 9:
+            return "Emergency response required. Notify environmental agencies."
+        else:
+            return "Disaster response needed. Evacuation may be necessary."
 
     def analyze_chat(self, message: str) -> str:
-        """Same chat handler"""
-        # ... (unchanged from your original code) ...
+        """Handle chat questions about pollution"""
+        prompt = f"""You are an environmental expert. Answer this question about river pollution: {message}
+        
+        Provide a concise, factual response in under 100 words."""
+        
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        outputs = self.model.generate(**inputs, max_new_tokens=150)
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        return response
 
 # Initialize with error handling
 try:
@@ -135,9 +204,38 @@ except Exception as e:
     analyzer = None
     model_status = f"❌ Model loading failed: {str(e)}"
 
-# Gradio Interface (unchanged layout from your original)
+# Gradio Interface
 css = """
-/* [Keep your existing CSS] */
+.header {
+    text-align: center;
+    max-width: 800px;
+    margin: auto;
+}
+.header img {
+    max-width: 100%;
+}
+.side-by-side {
+    display: flex;
+    flex-wrap: wrap;
+    gap: 20px;
+}
+.left-panel, .right-panel {
+    flex: 1;
+    min-width: 300px;
+}
+.analysis-box {
+    border: 1px solid #e0e0e0;
+    border-radius: 8px;
+    padding: 15px;
+    margin-top: 15px;
+    background: #f9f9f9;
+}
+.chat-container {
+    border: 1px solid #e0e0e0;
+    border-radius: 8px;
+    padding: 15px;
+    background: #f9f9f9;
+}
 """
 
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
@@ -177,7 +275,23 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         outputs=analysis_output
     )
 
-    # [Keep all other UI event handlers unchanged]
+    def respond(message, chat_history):
+        if not analyzer:
+            return chat_history + [(message, "Models not loaded. Please try again later.")]
+        response = analyzer.analyze_chat(message)
+        return chat_history + [(message, response)]
+
+    chat_btn.click(
+        respond,
+        [chat_input, chatbot],
+        [chatbot],
+    )
+    chat_input.submit(
+        respond,
+        [chat_input, chatbot],
+        [chatbot],
+    )
+    clear_btn.click(lambda: None, None, chatbot, queue=False)
 
     # Update examples to use local files
     gr.Examples(
@@ -192,5 +306,5 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         label="Try example images:"
     )
 
-# Launch with queue for stability
-demo.queue(max_size=3).launch()
+# Launch with queue for stability and allowed paths
+demo.queue(max_size=3).launch(allowed_paths=["examples"])