QA_based_Entity_Extraction.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ieg-dhr/NLP-Course4Humanities_2024/blob/main/QA_based_Entity_Extraction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LbB9AwCpIwIe"
      },
      "source": [
        "# QA-based Entity Extraction\n",
        "\n",
        "Created by Sarah Oberbichler [![ORCID](data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHdpZHRoPSIyMCIgaGVpZ2h0PSIyMCIgdmlld0JveD0iMCAwIDIwIDIwIj4KICA8cmVjdCB3aWR0aD0iMjAiIGhlaWdodD0iMjAiIGZpbGw9IiNGRkZGRkYiLz4KICA8Y2lyY2xlIGN4PSIxMCIgY3k9IjEwIiByPSI5IiBmaWxsPSIjQThDRTNDIi8+CiAgPHRleHQgeD0iMTAiIHk9IjE1IiBmb250LWZhbWlseT0iQXJpYWwsIHNhbnMtc2VyaWYiIGZvbnQtc2l6ZT0iMTEiIGZvbnQtd2VpZ2h0PSJib2xkIiB0ZXh0LWFuY2hvcj0ibWlkZGxlIiBmaWxsPSIjRkZGRkZGIj5pRDwvdGV4dD4KPC9zdmc+)](https://orcid.org/0000-0002-1031-2759)\n",
        "\n",
        "\n",
        "QA-based Entity Extraction transforms entity recognition into a question-answering task. Instead of directly labeling words in text as entities, it asks specific questions like \"What companies are mentioned?\" or \"Who are the people in this text?\" The model then responds by extracting the relevant entities from the text as answers to these questions. This approach makes entity extraction more flexible and intuitive, as new entity types can be added simply by asking new questions, though it may be more computationally intensive than traditional sequence labeling methods.\n",
        "\n",
        "###NuExtract Lannguage Model\n",
        "\n",
        "For the NE extraction, we use the NuExtract model v1.5. NuExtract is trained on a private high-quality dataset for structured information extraction. It supports long documents and several languages (English, French, Spanish, German, Portuguese, and Italian)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CEUT-wadXyZC"
      },
      "source": [
        "## Importing the Dataset\n",
        "\n",
        "We import a dataset that contains single articles"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G2Wpvvub7IxW"
      },
      "outputs": [],
      "source": [
        "!git clone https://github.com/ieg-dhr/NLP-Course4Humanities_2024.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K7ZY-SsI7LyY"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "\n",
        "articles_df = pd.read_excel('/content/NLP-Course4Humanities_2024/datasets/earthquake_articles.xlsx')\n",
        "\n",
        "articles_df = articles_df[:20]\n",
        "articles_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bUDhZrrjZ27N"
      },
      "source": [
        "##Defining a Template for Information Extraction\n",
        "\n",
        "As an example, we extract information from earthquake reportings. In doing so, we want the model to distinguish between earthquake locations, dateline locations, extract the date of the earthquake, the magnitutes, the persons involved as well as causalities, damage and rescue effort of the earthquake."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m62f_xoGZ56j"
      },
      "outputs": [],
      "source": [
        "# Define a template for earthquake information extraction\n",
        "earthquake_template = \"\"\"{\n",
        "    \"Earthquake\": {\n",
        "        \"Earthquake Locations\": \"\",\n",
        "        \"Dateline Locations\": \"\",\n",
        "        \"Date\": \"\",\n",
        "        \"Magnitude\": \"\",\n",
        "        \"Persons_Involved\": [],\n",
        "        \"Casualties\": {\n",
        "            \"Fatalities\": \"\",\n",
        "            \"Injured\": \"\"\n",
        "        },\n",
        "        \"Damage\": {\n",
        "            \"Infrastructure Damage\": \"\",\n",
        "            \"Economic Impact\": \"\"\n",
        "        },\n",
        "        \"Rescue Efforts\": \"\"\n",
        "    }\n",
        "}\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Running the Model\n",
        "\n",
        "The code below extracts the named entities using the extraction template. The model output is per default a json format. We add the extracted entities to our dataframe that will be saved as excel file."
      ],
      "metadata": {
        "id": "r-cCyVHIf5Sb"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "trbFNYPU69Ed"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import torch\n",
        "import pandas as pd\n",
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "\n",
        "def predict_NuExtract(model, tokenizer, texts, template, batch_size=1, max_length=10_000, max_new_tokens=4_000):\n",
        "    \"\"\"\n",
        "    Extract structured information from texts using NuExtract model\n",
        "\n",
        "    :param model: Loaded NuExtract model\n",
        "    :param tokenizer: Corresponding tokenizer\n",
        "    :param texts: List of input texts to extract from\n",
        "    :param template: JSON template for structured extraction\n",
        "    :param batch_size: Number of texts to process in parallel\n",
        "    :param max_length: Maximum input length\n",
        "    :param max_new_tokens: Maximum tokens to generate in output\n",
        "    :return: List of extracted information\n",
        "    \"\"\"\n",
        "    template = json.dumps(json.loads(template), indent=4)\n",
        "    prompts = [f\"\"\"<|input|>\\n### Template:\\n{template}\\n### Text:\\n{text}\\n\\n<|output|>\"\"\" for text in texts]\n",
        "\n",
        "    outputs = []\n",
        "    with torch.no_grad():\n",
        "        for i in range(0, len(prompts), batch_size):\n",
        "            batch_prompts = prompts[i:i+batch_size]\n",
        "            batch_encodings = tokenizer(batch_prompts, return_tensors=\"pt\", truncation=True, padding=True, max_length=max_length).to(model.device)\n",
        "\n",
        "            pred_ids = model.generate(**batch_encodings, max_new_tokens=max_new_tokens)\n",
        "            outputs += tokenizer.batch_decode(pred_ids, skip_special_tokens=True)\n",
        "\n",
        "    return [output.split(\"<|output|>\")[1] for output in outputs]\n",
        "    print(outputs)\n",
        "\n",
        "# Load the NuExtract model\n",
        "model_name = \"numind/NuExtract-v1.5\"\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, trust_remote_code=True).to(device).eval()\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
        "\n",
        "\n",
        "# Filter out non-string or empty entries\n",
        "articles_df['extracted_article_clean'] = articles_df['extracted_article_clean'].astype(str)\n",
        "valid_texts = articles_df['extracted_article_clean'][articles_df['extracted_article_clean'].str.strip() != '']\n",
        "\n",
        "# Extract information in batches to manage memory\n",
        "batch_size = 10  # Adjust based on your available memory\n",
        "all_predictions = []\n",
        "\n",
        "for i in range(0, len(valid_texts), batch_size):\n",
        "    batch_texts = valid_texts[i:i+batch_size].tolist()\n",
        "    batch_predictions = predict_NuExtract(model, tokenizer, batch_texts, earthquake_template)\n",
        "    all_predictions.extend(batch_predictions)\n",
        "\n",
        "# Parse predictions and add to DataFrame\n",
        "articles_df['earthquake_extraction'] = pd.Series([None] * len(articles_df))\n",
        "articles_df.loc[valid_texts.index, 'earthquake_extraction'] = all_predictions\n",
        "\n",
        "# Optional: Flatten the JSON extraction for easier analysis\n",
        "def parse_earthquake_info(extraction):\n",
        "    try:\n",
        "        parsed = json.loads(extraction)\n",
        "        return parsed.get('Earthquake', {})\n",
        "    except:\n",
        "        return {}\n",
        "\n",
        "articles_df['earthquake_locations'] = articles_df['earthquake_extraction'].apply(lambda x: parse_earthquake_info(x).get('Earthquake Locations', ''))\n",
        "articles_df['dateline_locations'] = articles_df['earthquake_extraction'].apply(lambda x: parse_earthquake_info(x).get('Dateline Locations', ''))\n",
        "articles_df['date'] = articles_df['earthquake_extraction'].apply(lambda x: parse_earthquake_info(x).get('Date', ''))\n",
        "articles_df['magnitude'] = articles_df['earthquake_extraction'].apply(lambda x: parse_earthquake_info(x).get('Magnitude', ''))\n",
        "articles_df['persons_involved'] = articles_df['earthquake_extraction'].apply(lambda x: parse_earthquake_info(x).get('Persons_Involved', []))\n",
        "articles_df['causalaties'] = articles_df['earthquake_extraction'].apply(lambda x: parse_earthquake_info(x).get('Injured', {}).get('Fatalities', ''))\n",
        "articles_df['infrastructure_damage'] = articles_df['earthquake_extraction'].apply(lambda x: parse_earthquake_info(x).get('Infrastructure Damage', ''))\n",
        "articles_df['economic_impact'] = articles_df['earthquake_extraction'].apply(lambda x: parse_earthquake_info(x).get('Economic Impact', ''))\n",
        "articles_df['rescue_effort'] = articles_df['earthquake_extraction'].apply(lambda x: parse_earthquake_info(x).get('Rescue Efforts', ''))\n",
        "\n",
        "\n",
        "# Optional: Save results\n",
        "articles_df.to_excel('earthquake_extractions.xlsx', index=False)\n",
        "\n",
        "\n",
        "articles_df"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Visualization Example - Creating a Map with Earthquake locations\n",
        "\n",
        "We first use the geopy library to process geographic locations and add their corresponding coordinates (latitude and longitude) to a pandas DataFrame. It includes a GeocodingService class that interfaces with the Nominatim geocoding API, implementing rate-limiting, retries with exponential backoff, and error handling to ensure robust geocoding.\n",
        "\n",
        "We further use the folium library to create an interactive map with markers for locations provided in a pandas DataFrame. Finally, the map is created and displayed, providing a visual representation of the geographic data."
      ],
      "metadata": {
        "id": "8B-_hff0g4IO"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ygivtCLFk4KD"
      },
      "outputs": [],
      "source": [
        "from geopy.geocoders import Nominatim\n",
        "from geopy.exc import GeocoderTimedOut, GeocoderServiceError\n",
        "import pandas as pd\n",
        "import time\n",
        "from typing import List, Tuple, Optional\n",
        "import random\n",
        "\n",
        "class GeocodingService:\n",
        "    def __init__(self, user_agent: str = None, timeout: int = 10, rate_limit: float = 1.1):\n",
        "        \"\"\"\n",
        "        Initialize the geocoding service with proper configuration.\n",
        "\n",
        "        Args:\n",
        "            user_agent: Custom user agent string (default: generated)\n",
        "            timeout: Timeout for requests in seconds\n",
        "            rate_limit: Time to wait between requests in seconds\n",
        "        \"\"\"\n",
        "        if user_agent is None:\n",
        "            user_agent = f\"python_geocoding_script_{random.randint(1000, 9999)}\"\n",
        "\n",
        "        self.geolocator = Nominatim(\n",
        "            user_agent=user_agent,\n",
        "            timeout=timeout\n",
        "        )\n",
        "        self.rate_limit = rate_limit\n",
        "        self.last_request = 0\n",
        "\n",
        "    def _rate_limit_wait(self):\n",
        "        \"\"\"Implement rate limiting between requests\"\"\"\n",
        "        current_time = time.time()\n",
        "        time_since_last = current_time - self.last_request\n",
        "        if time_since_last < self.rate_limit:\n",
        "            time.sleep(self.rate_limit - time_since_last)\n",
        "        self.last_request = time.time()\n",
        "\n",
        "    def geocode_location(self, location: str, max_retries: int = 3) -> Optional[Tuple[float, float]]:\n",
        "        \"\"\"\n",
        "        Geocode a single location with retries.\n",
        "\n",
        "        Args:\n",
        "            location: Location string to geocode\n",
        "            max_retries: Maximum number of retry attempts\n",
        "\n",
        "        Returns:\n",
        "            Tuple of (latitude, longitude) or None if geocoding fails\n",
        "        \"\"\"\n",
        "        for attempt in range(max_retries):\n",
        "            try:\n",
        "                self._rate_limit_wait()\n",
        "                location_data = self.geolocator.geocode(location)\n",
        "                if location_data:\n",
        "                    return (location_data.latitude, location_data.longitude)\n",
        "                return None\n",
        "            except (GeocoderTimedOut, GeocoderServiceError) as e:\n",
        "                if attempt == max_retries - 1:\n",
        "                    print(f\"Failed to geocode '{location}' after {max_retries} attempts: {e}\")\n",
        "                    return None\n",
        "                time.sleep(2 ** attempt)  # Exponential backoff\n",
        "            except Exception as e:\n",
        "                print(f\"Error geocoding '{location}': {e}\")\n",
        "                return None\n",
        "        return None\n",
        "\n",
        "    def process_locations(self, locations: str) -> List[Optional[Tuple[float, float]]]:\n",
        "        \"\"\"\n",
        "        Process a comma-separated string of locations.\n",
        "\n",
        "        Args:\n",
        "            locations: Comma-separated string of location names\n",
        "\n",
        "        Returns:\n",
        "            List of coordinate tuples or None for failed geocoding\n",
        "        \"\"\"\n",
        "        if pd.isna(locations) or not locations:\n",
        "            return []\n",
        "\n",
        "        location_list = [loc.strip() for loc in locations.split(',')]\n",
        "        return [self.geocode_location(loc) for loc in location_list]\n",
        "\n",
        "def geolocate_places(df: pd.DataFrame,\n",
        "                    places_column: str = 'places',\n",
        "                    user_agent: str = None) -> pd.DataFrame:\n",
        "    \"\"\"\n",
        "    Add coordinates to a DataFrame based on location names.\n",
        "\n",
        "    Args:\n",
        "        df: Input DataFrame\n",
        "        places_column: Name of the column containing comma-separated location strings\n",
        "        user_agent: Custom user agent string\n",
        "\n",
        "    Returns:\n",
        "        DataFrame with added 'coordinates' column\n",
        "    \"\"\"\n",
        "    geocoder = GeocodingService(user_agent=user_agent)\n",
        "\n",
        "    # Create a copy to avoid modifying the original DataFrame\n",
        "    result_df = df.copy()\n",
        "\n",
        "    # Process locations\n",
        "    result_df['coordinates'] = result_df[places_column].apply(geocoder.process_locations)\n",
        "\n",
        "    return result_df\n",
        "\n",
        "# Main execution\n",
        "if __name__ == \"__main__\":\n",
        "    # Assuming articles_df is your DataFrame with a 'places' column\n",
        "    # Apply geocoding to the articles DataFrame\n",
        "    articles_df_with_coords = geolocate_places(\n",
        "        articles_df,\n",
        "        places_column='earthquake_locations',\n",
        "        user_agent='article_geocoding_service_v1.0'\n",
        "    )\n",
        "\n",
        "    # Update the original DataFrame with the new coordinates\n",
        "    articles_df['coordinates'] = articles_df_with_coords['coordinates']\n",
        "\n",
        "    # Display the results\n",
        "    print(\"\\nSample of geocoded locations:\")\n",
        "    print(articles_df[['earthquake_locations', 'coordinates']].head())\n",
        "\n",
        "    # Optional: Display some statistics\n",
        "    total_locations = len(articles_df)\n",
        "    successful_geocodes = articles_df['coordinates'].apply(lambda x: len([c for c in x if c is not None])).sum()\n",
        "    failed_geocodes = articles_df['coordinates'].apply(lambda x: len([c for c in x if c is None])).sum()\n",
        "\n",
        "    print(f\"\\nGeocoding Statistics:\")\n",
        "    print(f\"Total locations processed: {total_locations}\")\n",
        "    print(f\"Successfully geocoded: {successful_geocodes}\")\n",
        "    print(f\"Failed to geocode: {failed_geocodes}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JbsClLRclZHY"
      },
      "outputs": [],
      "source": [
        "import folium\n",
        "from folium import plugins\n",
        "import pandas as pd\n",
        "from typing import List, Tuple, Optional\n",
        "from IPython.display import display\n",
        "\n",
        "def create_location_map(df: pd.DataFrame,\n",
        "                       coordinates_col: str = 'coordinates',\n",
        "                       places_col: str = 'earthquake_locations',\n",
        "                       title_col: Optional[str] = None) -> folium.Map:\n",
        "    \"\"\"\n",
        "    Create an interactive map with individual markers for all earthquake locations.\n",
        "\n",
        "    Args:\n",
        "        df: DataFrame containing coordinates and earthquake locations\n",
        "        coordinates_col: Name of column containing coordinates\n",
        "        places_col: Name of column containing earthquake location names\n",
        "        title_col: Optional column name for additional marker information\n",
        "\n",
        "    Returns:\n",
        "        folium.Map object with all locations marked individually\n",
        "    \"\"\"\n",
        "    # Initialize the map\n",
        "    m = folium.Map(location=[0, 0], zoom_start=2)\n",
        "\n",
        "    # Keep track of all valid coordinates for setting bounds\n",
        "    all_coords = []\n",
        "\n",
        "    # Process each row in the DataFrame\n",
        "    for idx, row in df.iterrows():\n",
        "        coordinates = row[coordinates_col]\n",
        "        places = row[places_col].split(',') if pd.notna(row[places_col]) else []\n",
        "        title = row[title_col] if title_col and pd.notna(row[title_col]) else None\n",
        "\n",
        "        # Skip if no coordinates\n",
        "        if not coordinates:\n",
        "            continue\n",
        "\n",
        "        # Add individual markers for each location\n",
        "        for i, (coord, place) in enumerate(zip(coordinates, places)):\n",
        "            if coord is not None:  # Skip None coordinates\n",
        "                lat, lon = coord\n",
        "                place_name = place.strip()\n",
        "\n",
        "                # Create popup content\n",
        "                popup_content = f\"<b>{place_name}</b>\"\n",
        "                if title:\n",
        "                    popup_content += f\"<br>{title}\"\n",
        "\n",
        "                # Add marker directly to the map (not in a cluster)\n",
        "                folium.Marker(\n",
        "                    location=[lat, lon],\n",
        "                    popup=folium.Popup(popup_content, max_width=300),\n",
        "                    tooltip=place_name,\n",
        "                    #icon=folium.Icon(color='red', icon='info-sign')\n",
        "                ).add_to(m)\n",
        "\n",
        "                all_coords.append([lat, lon])\n",
        "\n",
        "    # If we have coordinates, fit the map bounds to include all points\n",
        "    if all_coords:\n",
        "        m.fit_bounds(all_coords)\n",
        "\n",
        "    return m\n",
        "\n",
        "# Create and display the map\n",
        "map_obj = create_location_map(articles_df)\n",
        "display(map_obj)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "A100",
      "machine_shape": "hm",
      "provenance": [],
      "authorship_tag": "ABX9TyP6zRXdYDmZCSGy4qC8xbqH",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}