Update app.py

app.py

@@ -1,18 +1,15 @@
 import gradio as gr
 from Bio import Entrez
 from transformers import pipeline
-import spacy
 import os  # For environment variables and file paths
 
 # ---------------------------- Configuration ----------------------------
 ENTREZ_EMAIL = os.environ.get("ENTREZ_EMAIL", "[email protected]")  # Use environment variable, default fallback
 HUGGINGFACE_API_TOKEN = os.environ.get("HUGGINGFACE_API_TOKEN", "HUGGINGFACE_API_TOKEN") # Use environment variable, default fallback
 SUMMARIZATION_MODEL = "facebook/bart-large-cnn"
-SPACY_MODEL = "en_core_web_sm"
 
 # ---------------------------- Global Variables ----------------------------
 summarizer = None
-nlp = None
 initialization_status = "Initializing..."  # Track initialization state
 
 # ---------------------------- Helper Functions ----------------------------
@@ -26,39 +23,6 @@ def log_error(message: str):
     except:
         print("Couldn't write to error log file.")  #If logging fails, still print to console
 
-# ---------------------------- Language Model Loading ----------------------------
-
-def load_spacy_model(model_name="en_core_web_sm"):
-    """Loads the SpaCy language model, downloading it if necessary."""
-    global initialization_status  # To update the initialization status
-
-    try:
-        print(f"Attempting to load SpaCy model '{model_name}'...")
-        nlp_model = spacy.load(model_name)
-        print(f"Successfully loaded SpaCy model '{model_name}'.")
-        initialization_status += f"\nSpaCy model '{model_name}' loaded."
-        return nlp_model
-    except OSError:
-        print(f"SpaCy model '{model_name}' not found. Downloading...")
-        initialization_status += f"\nSpaCy model '{model_name}' not found. Downloading..."
-        try:
-            import subprocess
-            subprocess.check_call(["python", "-m", "spacy", "download", model_name])
-            nlp_model = spacy.load(model_name)
-            print(f"Successfully loaded SpaCy model '{model_name}' after downloading.")
-            initialization_status += f"\nSuccessfully loaded SpaCy model '{model_name}' after downloading."
-            return nlp_model
-
-        except Exception as e:
-            log_error(f"Failed to download or load SpaCy model '{model_name}': {e}")
-            initialization_status += f"\nFailed to download or load SpaCy model '{model_name}': {e}"
-            return None  # Indicate failure
-
-    except Exception as e:
-        log_error(f"Error loading SpaCy model '{model_name}': {e}")
-        initialization_status += f"\nError loading SpaCy model '{model_name}': {e}"
-        return None
-
 # ---------------------------- Tool Functions ----------------------------
 
 def search_pubmed(query: str) -> list:
@@ -95,27 +59,13 @@ def summarize_abstract(abstract: str) -> str:
     try:
         # Check if the abstract is empty or too short
         if not abstract or len(abstract.strip()) < 50:
-            return "Abstract too short to summarize."
+            return "Abstract too short to summarize. A more detailed abstract was not found."
 
         summary = summarizer(abstract, max_length=130, min_length=30, do_sample=False)[0]['summary_text']
         return summary
     except Exception as e:
         log_error(f"Summarization error: {e}")
-        return f"Error during summarization: {e}"
-
-def extract_entities(text: str) -> list:
-    """Extracts entities (simplified) using SpaCy."""
-    global nlp
-    if nlp is None:
-        log_error("SpaCy model not initialized.")
-        return "SpaCy model not initialized. Check initialization status."
-    try:
-        doc = nlp(text)
-        entities = [(ent.text, ent.label_) for ent in doc.ents]
-        return entities
-    except Exception as e:
-        log_error(f"Entity extraction error: {e}")
-        return [f"Error during entity extraction: {e}"]
+        return f"Error during summarization: Failed to generate concise summary with the current model."
 
 # ---------------------------- Agent Function ----------------------------
 
@@ -129,10 +79,15 @@ def medai_agent(query: str) -> str:
             abstract = fetch_abstract(article_id)
             if "Error" not in abstract:
                 summary = summarize_abstract(abstract)
-                entities = extract_entities(abstract)
-                results.append(f"**Article ID:** {article_id}\n\n**Summary:** {summary}\n\n**Entities:** {entities}\n\n---\n")
+                results.append(f"<div class='article'>\n"
+                               f"  <h3 class='article-id'>Article ID: {article_id}</h3>\n"
+                               f"  <p class='summary'><strong>Summary:</strong> {summary}</p>\n"
+                               f"</div>\n")
             else:
-                results.append(f"Error processing article {article_id}: {abstract}\n\n---\n")
+                results.append(f"<div class='article error'>\n"
+                               f"  <h3 class='article-id'>Article ID: {article_id}</h3>\n"
+                               f"  <p class='error-message'>Error processing article: {abstract}</p>\n"
+                               f"</div>\n")
         return "\n".join(results)
     else:
         return f"No articles found or error occurred: {article_ids}"
@@ -140,48 +95,59 @@ def medai_agent(query: str) -> str:
 # ---------------------------- Initialization and Setup ----------------------------
 
 def setup():
-    """Initializes the summarization model and SpaCy model."""
-    global summarizer, nlp, initialization_status
+    """Initializes the summarization model."""
+    global summarizer, initialization_status
     initialization_status = "Initializing..."
     try:
         print("Initializing summarization pipeline...")
         initialization_status += "\nInitializing summarization pipeline..."
         summarizer = pipeline("summarization", model=SUMMARIZATION_MODEL, token=HUGGINGFACE_API_TOKEN)
         print("Summarization pipeline initialized.")
-        initialization_status += "\nSummarization pipeline initialized."
-
-        print("Loading SpaCy model...")
-        initialization_status += "\nLoading SpaCy model..."
-        global nlp
-        nlp = load_spacy_model()  # Call the SpaCy loading function.
-        if nlp is None:
-            initialization_status += "\nSpaCy model failed to load. Check the error log."
-            return initialization_status
-
-        print("SpaCy model loaded.")
-        initialization_status += "\nSpaCy model loaded."
-
-        initialization_status = "MedAI Agent initialized successfully!"
-        return initialization_status  # Return the status message
+        initialization_status += f"\nSummarization pipeline initialized. Model {SUMMARIZATION_MODEL} loaded and ready."
+        return initialization_status
     except Exception as e:
         initialization_status = f"Initialization error: {e}"
         log_error(initialization_status)
-        return initialization_status  # Return the error message
+        return initialization_status
 
 # ---------------------------- Gradio Interface ----------------------------
 
 def launch_gradio():
     """Launches the Gradio interface."""
-    global initialization_status  # Allows the function to modify global variable
-    with gr.Blocks() as iface:
+    global initialization_status
+
+    # CSS to style the article output
+    css = """
+    .article {
+      border: 1px solid #ddd;
+      margin-bottom: 10px;
+      padding: 10px;
+      border-radius: 5px;
+    }
+    .article.error {
+      border-color: #f00;
+    }
+    .article-id {
+      font-size: 1.2em;
+      margin-bottom: 5px;
+    }
+    .summary {
+      font-style: italic;
+    }
+    .error-message {
+      color: #f00;
+    }
+    """
+
+    with gr.Blocks(css=css) as iface:
         gr.Markdown("# MedAI: Medical Literature Review and Summarization")
-        status_display = gr.Textbox(value=initialization_status, interactive=False)  # Displays initialization status
+        status_display = gr.Textbox(value=initialization_status, interactive=False)
         query_input = gr.Textbox(lines=3, placeholder="Enter your medical query (e.g., 'new treatments for diabetes')...")
         submit_button = gr.Button("Submit")
-        output_results = gr.Markdown()
+        output_results = gr.HTML()  # Use HTML for formatted output
 
         submit_button.click(medai_agent, inputs=query_input, outputs=output_results)
-        status_display.value = setup()  # Set the status after running setup
+        status_display.value = setup()
 
     iface.launch()