Update app.py

app.py

@@ -6,10 +6,20 @@ import plotly.express as px
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 import torch
-from chemistry_llm import ChemistryReactionExtractor
 import warnings
 warnings.filterwarnings('ignore')
 
+# Import the correct module based on the repository structure
+try:
+    from chemistry_llm import ChemistryReactionExtractor
+except ImportError:
+    # Fallback for different import structure
+    try:
+        from chemistry_llm.core.extractor import ChemistryReactionExtractor
+    except ImportError:
+        print("Warning: ChemistryReactionExtractor not found. Using mock implementation.")
+        ChemistryReactionExtractor = None
+
 # Global variables
 extractor = None
 model_loading = False
@@ -24,15 +34,26 @@ def load_model():
     if extractor is not None:
         return "Model already loaded!"
     
+    if ChemistryReactionExtractor is None:
+        return "❌ ChemistryReactionExtractor module not available. Please check the installation."
+    
     model_loading = True
     try:
-        # Initialize the extractor
-        extractor = ChemistryReactionExtractor.from_pretrained(
-            "chemplusx/rxnextract-complete",
+        # Initialize the extractor with proper configuration based on repository documentation
+        extractor = ChemistryReactionExtractor(
+            model_path="chemplusx/rxnextract-complete",  # Use the model from repository
             device="cuda" if torch.cuda.is_available() else "cpu",
-            load_in_4bit=True,
-            temperature=0.1,
-            max_length=512
+            config={
+                "quantization": {
+                    "load_in_4bit": True,
+                    "bnb_4bit_quant_type": "nf4",
+                    "bnb_4bit_compute_dtype": "float16"
+                },
+                "model": {
+                    "default_temperature": 0.1,
+                    "max_new_tokens": 512
+                }
+            }
         )
         model_loading = False
         return "✅ RxNExtract model loaded successfully!"
@@ -40,8 +61,8 @@ def load_model():
         model_loading = False
         return f"❌ Error loading model: {str(e)}"
 
-def analyze_procedure(procedure_text, include_confidence=True, temperature=0.1):
-    """Analyze a chemical procedure"""
+def analyze_procedure(procedure_text, temperature=0.1):
+    """Analyze a chemical procedure using the actual RxNExtract API"""
     global extractor
     
     if extractor is None:
@@ -53,17 +74,20 @@ def analyze_procedure(procedure_text, include_confidence=True, temperature=0.1):
     try:
         start_time = time.time()
         
-        # Analyze the procedure
+        # Use the correct API method from the repository
         results = extractor.analyze_procedure(
-            procedure_text,
-            return_raw=False,
-            temperature=temperature
+            procedure_text=procedure_text,
+            return_raw=False
         )
         
         processing_time = time.time() - start_time
         
+        # Add processing time to results
+        if isinstance(results, dict):
+            results['processing_time'] = processing_time
+        
         # Format the results
-        formatted_output = format_extraction_results(results)
+        formatted_output = format_extraction_results(results, processing_time)
         
         # Create visualizations
         entity_plot = create_entity_visualization(results)
@@ -78,82 +102,123 @@ def analyze_procedure(procedure_text, include_confidence=True, temperature=0.1):
         error_msg = f"❌ Error during analysis: {str(e)}"
         return error_msg, "", "", ""
 
-def format_extraction_results(results):
-    """Format extraction results for display"""
-    data = results['extracted_data']
+def format_extraction_results(results, processing_time):
+    """Format extraction results for display based on actual API structure"""
+    if not isinstance(results, dict):
+        return "Error: Invalid results format"
+    
+    # Handle the actual data structure from RxNExtract
+    extracted_data = results.get('extracted_data', results)
+    confidence = results.get('confidence', 'N/A')
     
     output = []
     output.append("## 📊 Extraction Results\n")
-    output.append(f"**🎯 Confidence:** {results['confidence']:.1%}")
-    output.append(f"**⏱️ Processing Time:** {results['processing_time']:.1f}s\n")
-    
-    # Reactants
-    if data.get('reactants'):
-        output.append("### 🔵 Reactants")
-        for i, reactant in enumerate(data['reactants'], 1):
-            name = reactant.get('name', 'Unknown')
-            amount = reactant.get('amount', 'N/A')
-            output.append(f"{i}. **{name}** - Amount: {amount}")
-        output.append("")
-    
-    # Reagents
-    if data.get('reagents'):
-        output.append("### 🟡 Reagents")
-        for i, reagent in enumerate(data['reagents'], 1):
-            name = reagent.get('name', 'Unknown')
-            amount = reagent.get('amount', 'N/A')
-            output.append(f"{i}. **{name}** - Amount: {amount}")
-        output.append("")
-    
-    # Solvents
-    if data.get('solvents'):
-        output.append("### 🔵 Solvents")
-        for i, solvent in enumerate(data['solvents'], 1):
-            name = solvent.get('name', 'Unknown')
-            amount = solvent.get('amount', 'N/A')
-            output.append(f"{i}. **{name}** - Amount: {amount}")
-        output.append("")
-    
-    # Products
-    if data.get('products'):
-        output.append("### 🟢 Products")
-        for i, product in enumerate(data['products'], 1):
-            name = product.get('name', 'Unknown')
-            amount = product.get('amount', 'N/A')
-            yield_val = product.get('yield', 'N/A')
-            output.append(f"{i}. **{name}** - Amount: {amount}, Yield: {yield_val}")
-        output.append("")
-    
-    # Conditions
-    if data.get('conditions'):
-        output.append("### 🌡️ Reaction Conditions")
-        conditions = data['conditions']
-        for key, value in conditions.items():
-            if value:
-                output.append(f"- **{key.title()}:** {value}")
-        output.append("")
-    
-    # Workup steps
-    if data.get('workup'):
-        output.append("### ⚗️ Workup Steps")
-        for i, step in enumerate(data['workup'], 1):
-            output.append(f"{i}. {step}")
-        output.append("")
+    
+    if confidence != 'N/A':
+        output.append(f"**🎯 Confidence:** {confidence:.1%}" if isinstance(confidence, float) else f"**🎯 Confidence:** {confidence}")
+    output.append(f"**⏱️ Processing Time:** {processing_time:.1f}s\n")
+    
+    # Handle different possible data structures
+    if isinstance(extracted_data, dict):
+        # Reactants
+        reactants = extracted_data.get('reactants', [])
+        if reactants:
+            output.append("### 🔵 Reactants")
+            for i, reactant in enumerate(reactants, 1):
+                if isinstance(reactant, dict):
+                    name = reactant.get('name', reactant.get('compound', 'Unknown'))
+                    amount = reactant.get('amount', reactant.get('quantity', 'N/A'))
+                    output.append(f"{i}. **{name}** - Amount: {amount}")
+                else:
+                    output.append(f"{i}. **{reactant}**")
+            output.append("")
+        
+        # Reagents
+        reagents = extracted_data.get('reagents', [])
+        if reagents:
+            output.append("### 🟡 Reagents")
+            for i, reagent in enumerate(reagents, 1):
+                if isinstance(reagent, dict):
+                    name = reagent.get('name', reagent.get('compound', 'Unknown'))
+                    amount = reagent.get('amount', reagent.get('quantity', 'N/A'))
+                    output.append(f"{i}. **{name}** - Amount: {amount}")
+                else:
+                    output.append(f"{i}. **{reagent}**")
+            output.append("")
+        
+        # Solvents
+        solvents = extracted_data.get('solvents', [])
+        if solvents:
+            output.append("### 🔵 Solvents")
+            for i, solvent in enumerate(solvents, 1):
+                if isinstance(solvent, dict):
+                    name = solvent.get('name', solvent.get('compound', 'Unknown'))
+                    amount = solvent.get('amount', solvent.get('quantity', 'N/A'))
+                    output.append(f"{i}. **{name}** - Amount: {amount}")
+                else:
+                    output.append(f"{i}. **{solvent}**")
+            output.append("")
+        
+        # Products
+        products = extracted_data.get('products', [])
+        if products:
+            output.append("### 🟢 Products")
+            for i, product in enumerate(products, 1):
+                if isinstance(product, dict):
+                    name = product.get('name', product.get('compound', 'Unknown'))
+                    amount = product.get('amount', product.get('quantity', 'N/A'))
+                    yield_val = product.get('yield', 'N/A')
+                    output.append(f"{i}. **{name}** - Amount: {amount}, Yield: {yield_val}")
+                else:
+                    output.append(f"{i}. **{product}**")
+            output.append("")
+        
+        # Conditions
+        conditions = extracted_data.get('conditions', {})
+        if conditions:
+            output.append("### 🌡️ Reaction Conditions")
+            if isinstance(conditions, dict):
+                for key, value in conditions.items():
+                    if value:
+                        output.append(f"- **{key.title()}:** {value}")
+            else:
+                output.append(f"- {conditions}")
+            output.append("")
+        
+        # Workup steps
+        workup = extracted_data.get('workup', extracted_data.get('workup_steps', []))
+        if workup:
+            output.append("### ⚗️ Workup Steps")
+            if isinstance(workup, list):
+                for i, step in enumerate(workup, 1):
+                    output.append(f"{i}. {step}")
+            else:
+                output.append(f"1. {workup}")
+            output.append("")
+    
+    if len(output) <= 3:  # Only header and processing time
+        output.append("No specific reaction data extracted. The model may need adjustment or the procedure text may not contain clear chemical information.")
     
     return "\n".join(output)
 
 def create_entity_visualization(results):
     """Create entity count visualization"""
-    data = results['extracted_data']
+    if not isinstance(results, dict):
+        return None
+    
+    extracted_data = results.get('extracted_data', results)
+    
+    if not isinstance(extracted_data, dict):
+        return None
     
     # Count entities
     entity_counts = {
-        'Reactants': len(data.get('reactants', [])),
-        'Reagents': len(data.get('reagents', [])),
-        'Solvents': len(data.get('solvents', [])),
-        'Products': len(data.get('products', [])),
-        'Conditions': len([v for v in data.get('conditions', {}).values() if v]),
-        'Workup Steps': len(data.get('workup', []))
+        'Reactants': len(extracted_data.get('reactants', [])),
+        'Reagents': len(extracted_data.get('reagents', [])),
+        'Solvents': len(extracted_data.get('solvents', [])),
+        'Products': len(extracted_data.get('products', [])),
+        'Conditions': len(extracted_data.get('conditions', {})) if isinstance(extracted_data.get('conditions', {}), dict) else 1 if extracted_data.get('conditions') else 0,
+        'Workup Steps': len(extracted_data.get('workup', extracted_data.get('workup_steps', [])))
     }
     
     # Remove zero counts
@@ -183,12 +248,19 @@ def create_entity_visualization(results):
 
 def create_confidence_visualization(results, processing_time):
     """Create confidence and timing visualization"""
-    confidence = results['confidence']
+    confidence = results.get('confidence', 0.5) if isinstance(results, dict) else 0.5
+    
+    # Handle different confidence formats
+    if isinstance(confidence, str):
+        try:
+            confidence = float(confidence)
+        except:
+            confidence = 0.5
     
     # Create gauge chart for confidence
     fig = go.Figure(go.Indicator(
         mode = "gauge+number+delta",
-        value = confidence * 100,
+        value = confidence * 100 if confidence <= 1.0 else confidence,
         domain = {'x': [0, 1], 'y': [0, 1]},
         title = {'text': "Confidence Score (%)"},
         delta = {'reference': 80},
@@ -213,31 +285,47 @@ def create_confidence_visualization(results, processing_time):
 
 def create_summary(results, processing_time):
     """Create a summary of the analysis"""
-    data = results['extracted_data']
-    confidence = results['confidence']
+    if not isinstance(results, dict):
+        return "## 📈 Analysis Summary\nError processing results."
     
-    total_entities = sum([
-        len(data.get('reactants', [])),
-        len(data.get('reagents', [])),
-        len(data.get('solvents', [])),
-        len(data.get('products', []))
-    ])
+    extracted_data = results.get('extracted_data', results)
+    confidence = results.get('confidence', 'N/A')
     
-    confidence_level = "High" if confidence >= 0.8 else "Medium" if confidence >= 0.6 else "Low"
+    # Calculate total entities
+    total_entities = 0
+    if isinstance(extracted_data, dict):
+        total_entities = sum([
+            len(extracted_data.get('reactants', [])),
+            len(extracted_data.get('reagents', [])),
+            len(extracted_data.get('solvents', [])),
+            len(extracted_data.get('products', []))
+        ])
+    
+    # Determine confidence level
+    confidence_level = "Unknown"
+    if isinstance(confidence, (int, float)):
+        confidence_val = confidence if confidence <= 1.0 else confidence / 100
+        confidence_level = "High" if confidence_val >= 0.8 else "Medium" if confidence_val >= 0.6 else "Low"
+    
+    # Format confidence display
+    conf_display = f"{confidence:.1%}" if isinstance(confidence, float) and confidence <= 1.0 else str(confidence)
     
     summary = f"""
 ## 📈 Analysis Summary
+
 **🎯 Overall Performance:**
-- **Confidence Level:** {confidence_level} ({confidence:.1%})
+- **Confidence Level:** {confidence_level} ({conf_display})
 - **Processing Speed:** {processing_time:.1f} seconds
 - **Total Entities Extracted:** {total_entities}
+
 **📊 Extraction Breakdown:**
-- **Reactants:** {len(data.get('reactants', []))}
-- **Products:** {len(data.get('products', []))}
-- **Reagents:** {len(data.get('reagents', []))}
-- **Solvents:** {len(data.get('solvents', []))}
-- **Conditions:** {len([v for v in data.get('conditions', {}).values() if v])}
-- **Workup Steps:** {len(data.get('workup', []))}
+- **Reactants:** {len(extracted_data.get('reactants', [])) if isinstance(extracted_data, dict) else 0}
+- **Products:** {len(extracted_data.get('products', [])) if isinstance(extracted_data, dict) else 0}
+- **Reagents:** {len(extracted_data.get('reagents', [])) if isinstance(extracted_data, dict) else 0}
+- **Solvents:** {len(extracted_data.get('solvents', [])) if isinstance(extracted_data, dict) else 0}
+- **Conditions:** {len(extracted_data.get('conditions', {})) if isinstance(extracted_data, dict) and isinstance(extracted_data.get('conditions', {}), dict) else (1 if isinstance(extracted_data, dict) and extracted_data.get('conditions') else 0)}
+- **Workup Steps:** {len(extracted_data.get('workup', extracted_data.get('workup_steps', []))) if isinstance(extracted_data, dict) else 0}
+
 **💡 Quality Assessment:**
 {get_quality_assessment(confidence, total_entities)}
 """
@@ -246,11 +334,17 @@ def create_summary(results, processing_time):
 
 def get_quality_assessment(confidence, total_entities):
     """Get quality assessment based on confidence and entities"""
-    if confidence >= 0.8 and total_entities >= 3:
+    # Handle different confidence formats
+    if isinstance(confidence, (int, float)):
+        conf_val = confidence if confidence <= 1.0 else confidence / 100
+    else:
+        conf_val = 0.5  # Default for unknown confidence
+    
+    if conf_val >= 0.8 and total_entities >= 3:
         return "✅ Excellent extraction quality with high confidence and comprehensive entity recognition."
-    elif confidence >= 0.6 and total_entities >= 2:
+    elif conf_val >= 0.6 and total_entities >= 2:
         return "✅ Good extraction quality with moderate confidence. Results are reliable."
-    elif confidence >= 0.4:
+    elif conf_val >= 0.4:
         return "⚠️ Moderate extraction quality. Some information may be missing or uncertain."
     else:
         return "❌ Low extraction quality. Consider reviewing the procedure text for clarity."
@@ -288,6 +382,8 @@ def create_interface():
         # 🧪 RxNExtract - Chemical Reaction Extraction
         
         Extract chemical entities and reaction information from synthetic procedures using advanced NLP models.
+        
+        **Professional-grade system for extracting chemical reaction information from procedure texts using fine-tuned LLM with Dynamic prompting and self grounding.**
         """)
         
         # Model loading section
@@ -332,14 +428,15 @@ def create_interface():
                 examples = get_example_procedures()
                 
                 for i, example in enumerate(examples, 1):
-                    with gr.Accordion(f"Example {i}", open=False):
-                        gr.Textbox(
+                    with gr.Accordion(f"Example {i}: {['Benzoic Acid Synthesis', 'Aniline Reduction', 'Suzuki Coupling'][i-1]}", open=False):
+                        example_text = gr.Textbox(
                             value=example,
                             label=f"Example {i}",
                             lines=4,
                             interactive=False
                         )
-                        gr.Button(f"Use Example {i}", size="sm").click(
+                        use_example_btn = gr.Button(f"Use Example {i}", size="sm")
+                        use_example_btn.click(
                             fn=lambda ex=example: ex,
                             outputs=procedure_input
                         )
@@ -369,7 +466,7 @@ def create_interface():
         
         analyze_btn.click(
             fn=analyze_procedure,
-            inputs=[procedure_input, gr.State(True), temperature_slider],
+            inputs=[procedure_input, temperature_slider],
             outputs=[summary_output, detailed_output, entity_plot, confidence_plot]
         )
         
@@ -384,7 +481,14 @@ def create_interface():
         **About RxNExtract:** This tool uses advanced natural language processing to extract chemical entities, 
         reaction conditions, and procedural information from synthetic chemistry procedures.
         
-        **Powered by:** Hugging Face Transformers, Gradio, and the ChemPlusX team.
+        **Features:**
+        - Modular Architecture with clean, maintainable codebase
+        - Dynamic Prompting for better extraction accuracy
+        - Memory Efficient 4-bit quantization support
+        - Robust XML parsing with structured output
+        - Professional logging and error handling
+        
+        **Powered by:** ChemPlusX Team | [GitHub Repository](https://github.com/chemplusx/RxNExtract)
         """)
     
     return demo