--- a +++ b/aitrika/base/aitrika_base.py @@ -0,0 +1,131 @@ +from typing import Optional, Dict, List, Union +import pandas as pd +from aitrika.extractors.pubmed_extractor import PubMedExtractor +from aitrika.extractors.llm_extractor import LLMExtractor +from aitrika.utils.loader import loader +from aitrika.llm.base_llm import BaseLLM + + +class AItrikaBase: + def __init__(self, pubmed_id: Optional[str] = None): + self.pubmed_id = pubmed_id + self.pubmed_extractor = PubMedExtractor(pubmed_id) if pubmed_id else None + self.llm_extractor = LLMExtractor() + + @loader(text="Extracting PubMed ID") + def extract_pubmed_id(self) -> str: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_pubmed_id() + + @loader(text="Extracting title") + def extract_title(self) -> str: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_title() + + @loader(text="Extracting abstract") + def extract_abstract(self) -> str: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_abstract() + + @loader(text="Extracting other abstract") + def extract_other_abstract(self) -> str: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_other_abstract() + + @loader(text="Extracting genes") + def extract_genes(self, dataframe: bool = False) -> Union[pd.DataFrame, str]: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_genes(dataframe) + + @loader(text="Extracting diseases") + def extract_diseases(self, dataframe: bool = False) -> Union[pd.DataFrame, str]: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_diseases(dataframe) + + @loader(text="Extracting species") + def extract_species(self, dataframe: bool = False) -> Union[pd.DataFrame, str]: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_species(dataframe) + + @loader(text="Extracting chemicals") + def extract_chemicals(self, dataframe: bool = False) -> Union[pd.DataFrame, str]: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_chemicals(dataframe) + + @loader(text="Extracting mutations") + def extract_mutations(self, dataframe: bool = False) -> Union[pd.DataFrame, str]: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_mutations(dataframe) + + @loader(text="Extracting associations between genes and diseases") + def extract_associations( + self, dataframe: bool = False + ) -> Union[pd.DataFrame, List[Dict]]: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_associations(dataframe) + + @loader(text="Extracting authors") + def extract_authors(self) -> str: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_authors() + + @loader(text="Extracting journal") + def extract_journal(self) -> str: + if not self.pubmed_extractor: + raise ValueError("PubMed ID not set") + return self.pubmed_extractor.extract_journal() + + @loader(text="Extracting results") + def extract_results(self, llm: BaseLLM) -> str: + return self.llm_extractor.extract_results(llm) + + @loader(text="Extracting bibliography") + def extract_bibliography(self, llm: BaseLLM) -> str: + return self.llm_extractor.extract_bibliography(llm) + + @loader(text="Extracting methods") + def extract_methods(self, llm: BaseLLM) -> str: + return self.llm_extractor.extract_methods(llm) + + @loader(text="Extracting introduction") + def extract_introduction(self, llm: BaseLLM) -> str: + return self.llm_extractor.extract_introduction(llm) + + @loader(text="Extracting acknowledgements") + def extract_acknowledgements(self, llm: BaseLLM) -> str: + return self.llm_extractor.extract_acknowledgements(llm) + + @loader(text="Extracting paper results") + def extract_paper_results(self, llm: BaseLLM) -> str: + return self.llm_extractor.extract_paper_results(llm) + + @loader(text="Extracting effect sizes") + def extract_effect_sizes(self, llm: BaseLLM) -> str: + return self.llm_extractor.extract_effect_sizes(llm) + + @loader(text="Extracting number of participants") + def extract_number_of_participants(self, llm: BaseLLM) -> str: + return self.llm_extractor.extract_number_of_participants(llm) + + @loader(text="Extracting characteristics of participants") + def extract_characteristics_of_participants(self, llm: BaseLLM) -> str: + return self.llm_extractor.extract_characteristics_of_participants(llm) + + @loader(text="Extracting interventions") + def extract_interventions(self, llm: BaseLLM) -> str: + return self.llm_extractor.extract_interventions(llm) + + @loader(text="Extracting outcomes") + def extract_outcomes(self, llm: BaseLLM) -> str: + return self.llm_extractor.extract_outcomes(llm)