dataset.py

import requests
import re
import json
import argparse

# Wikipedia & Wikidata API URLs
WIKI_API_URL = "https://pt.wikipedia.org/w/api.php"
WIKIDATA_API_URL = "https://www.wikidata.org/wiki/Special:EntityData/{}.json"
DATA_API_URL = "https://www.wikidata.org/w/api.php"

HEADERS = {
    "User-Agent": "MyWikidataTool/1.0 (https://github.com/yourusername/yourrepo; you@example.com)"
}

def translate_relation(input_string):
    mapping_dict = {
        "P31": "instância de",
        "P106": "ocupação",
        "P19": "local de nascimento",
        "P279": "subclasse de",
        "P17": "país",
        "P159": "sede",
        "P921": "tema(s) principal(is)",
        "P50": "autor",
        "P105": "categoria taxonómica",
        "P641": "desporto",
        "P136": "género",
        "P140": "religião",
        "P22": "pai",
        "P206": "banhado por",
        "P86":"compositor",
        "P272": "empresa de produção",
        "P178": "desenvolvedor",
        "P57": "realizador",
    }

    return mapping_dict.get(input_string, None)




# Function to get Wikipedia page title for an entity (from Wikidata ID)

def get_wikipedia_title(wikidata_id: str) -> str | None:
    url = WIKIDATA_API_URL.format(wikidata_id)
    try:
        response = requests.get(url, headers=HEADERS)
        response.raise_for_status()
        data = response.json()
    except requests.RequestException:
        return None
    except ValueError:
        return None

    try:
        return data["entities"][wikidata_id]["sitelinks"]["ptwiki"]["title"]
    except KeyError:
        return None  # No Wikipedia page found


def get_wikidata_label(wikidata_id):
    headers = {
        'User-Agent': 'MyWikiBot/1.0 (https://github.com/yourusername/yourrepo; myemail@example.com)'
    }

    params = {
        "action": "wbgetentities",
        "ids": wikidata_id,
        "format": "json",
        "props": "labels",
        "languages": "pt|en"
    }

    try:
        response = requests.get(DATA_API_URL, params=params, headers=headers)
        response.raise_for_status()  # Raises exception for 4xx/5xx status codes
        
        data = response.json()
        
        # Check if the entity exists and has labels
        if ("entities" in data and 
            wikidata_id in data["entities"] and 
            "labels" in data["entities"][wikidata_id]):
            
            labels = data["entities"][wikidata_id]["labels"]
            
            # Try Portuguese first, then English, then any available language
            if "pt" in labels:
                return labels["pt"].get("value")
            elif "en" in labels:
                return labels["en"].get("value")
            elif labels:  # If there are any labels, return the first one
                first_label = next(iter(labels.values()))
                return first_label.get("value")
        
        return None  # No labels found
        
    except requests.exceptions.RequestException as e:
        print(f"Request error for {wikidata_id}: {e}")
        return None
    except ValueError as e:  # JSON decoding error
        print(f"JSON decode error for {wikidata_id}: {e}")
        return None
    except Exception as e:
        print(f"Unexpected error for {wikidata_id}: {e}")
        return None




# Function to fetch Wikipedia page content and clean it
def get_wikipedia_text(title: str) -> str | None:
    params = {
        "action": "query",
        "format": "json",
        "prop": "extracts",
        "exintro": True,  # Only get the introduction
        "titles": title
    }

    try:
        response = requests.get(WIKI_API_URL, params=params, headers=HEADERS)
        response.raise_for_status()
        data = response.json()
    except requests.RequestException:
        return None
    except ValueError:
        return None


    pages = data["query"]["pages"]
    for page_id in pages:
        if "extract" in pages[page_id]:
            # Use BeautifulSoup to clean the HTML and extract text
            #soup = BeautifulSoup(pages[page_id]["extract"], "lxml")
            #return soup.get_text()
            return pages[page_id]["extract"]
    return None



def extract_sentences(text, entity1, entity2):
    """
    Extracts the first sentence where both entities appear together.

    Args:
        text (str): The input text to search in.
        entity1 (str): The first entity to search for.
        entity2 (str): The second entity to search for.

    Returns:
        str: The first sentence where both entities appear, or an empty array if no match is found.
    """
    # Split the text into sentences
    sentences = re.split(r"(?<=[.!?])(?:\s|\n|<)", text)

    # Use lower() for case-insensitive comparison
    entity1_lower = entity1.lower()
    entity2_lower = entity2.lower()

    # Create word-boundary regex patterns for stricter matching
    pattern1 = re.compile(rf"\b{re.escape(entity1_lower)}\b", re.IGNORECASE)
    pattern2 = re.compile(rf"\b{re.escape(entity2_lower)}\b", re.IGNORECASE)

    matched_sentences = []
    # Iterate through sentences and find the first match
    for sentence in sentences:
        sentence_lower = sentence.lower()

        match1 = pattern1.search(sentence_lower)
        match2 = pattern2.search(sentence_lower)

        if match1 and match2:
            
            clean_text = re.sub(r"<[^>]+>", "", sentence)  # Remove any HTML tags
           
            matched_sentences.append(clean_text.strip())

            return matched_sentences


    # Return an empty array if no match is found
    return matched_sentences


# Function to process a Wikidata triple and get Wikipedia sentences
def process_wikidata_triple(subject_id, relation, object_id, object_label, search_obj = 0):
    title = get_wikipedia_title(subject_id)

    if not title:
        print(f"❌ Sem página da Wikipédia para {subject_id}")
        return None

    text = get_wikipedia_text(title)
    if not text:
        print(f"❌ Nenhum texto encontrado para {title}")
        return None

    
    matched_sentences = extract_sentences(text, title, object_label)


    if len(matched_sentences) == 0:
        if search_obj:
            title2 = get_wikipedia_title(object_id)
            if title2 == None:
                print(f"❌ Sem página da Wikipédia para {object_id}")
                return None

            text2 = get_wikipedia_text(title2)
            if text2 == None:
                print(f"❌ Nenhum texto encontrado para {title2}")
                return None

           
            subject_label = get_wikidata_label(subject_id)
            if subject_label == None:
              return None

            matched_sentences2 = extract_sentences(text2, title2, subject_label)

            if len(matched_sentences2) == 0:
              return None
            else:
              return {"sentence": matched_sentences2[0], "subject": subject_label, "relation": relation, "object": title2}

        else:
            return None
    else:
        return {"sentence": matched_sentences[0], "subject": title, "relation": relation, "object": object_label}
        


if __name__ == "__main__":


    parser = argparse.ArgumentParser(description="dataset creation")
    parser.add_argument("input_file", type=str, help="Path to the input triples csv file.")
     
    parser.add_argument("-o", "--output_file", type=str, default="dataset.jsonl", help="output dataset file path")
    parser.add_argument("-f", "--filtered_file", type=str, default="filtered_triples.csv", help="path to the triples used on the dataset")


    args = parser.parse_args()


 

    
    print("Starting the Wikidata and Wikipedia connection")

    valid_lines = []

    with open(args.input_file, 'r') as infile, open(args.output_file, "w", encoding="utf-8") as f:
        for line in infile:
            subject_id, relation_id, object_id = line.strip().split('\t')  # Split the line into components
            object_label = get_wikidata_label(object_id)
            translated = translate_relation(relation_id)
            relation = translated if translated is not None else relation_id

            if object_label == None:
                continue

            sentences = process_wikidata_triple(subject_id, relation, object_id, object_label,1)
            
            if sentences == None:
                continue


            valid_lines.append(line)
            json.dump(sentences, f, ensure_ascii=False)
            f.write('\n')
            

    print("Creating the new filtered triples file")
    with open(args.filtered_file, 'w', encoding='utf-8') as outfile:
        outfile.writelines(valid_lines)