H commited on
Commit
c499f9e
·
1 Parent(s): 486f207

Add component pubmed (#1586)

Browse files

### What problem does this PR solve?

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

graph/component/__init__.py CHANGED
@@ -12,6 +12,7 @@ from .keyword import KeywordExtract, KeywordExtractParam
12
  from .baidu import Baidu, BaiduParam
13
  from .duckduckgo import DuckDuckGo, DuckDuckGoParam
14
  from .wikipedia import Wikipedia, WikipediaParam
 
15
 
16
 
17
  def component_class(class_name):
 
12
  from .baidu import Baidu, BaiduParam
13
  from .duckduckgo import DuckDuckGo, DuckDuckGoParam
14
  from .wikipedia import Wikipedia, WikipediaParam
15
+ from .pubmed import PubMed, PubMedParam
16
 
17
 
18
  def component_class(class_name):
graph/component/pubmed.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #
2
+ # Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ #
16
+ import random
17
+ from abc import ABC
18
+ from functools import partial
19
+ from Bio import Entrez
20
+ import pandas as pd
21
+ import xml.etree.ElementTree as ET
22
+ from graph.settings import DEBUG
23
+ from graph.component.base import ComponentBase, ComponentParamBase
24
+
25
+
26
+ class PubMedParam(ComponentParamBase):
27
+ """
28
+ Define the PubMed component parameters.
29
+ """
30
+
31
+ def __init__(self):
32
+ super().__init__()
33
+ self.top_n = 5
34
+ self.email = "[email protected]"
35
+
36
+ def check(self):
37
+ self.check_positive_integer(self.top_n, "Top N")
38
+
39
+
40
+ class PubMed(ComponentBase, ABC):
41
+ component_name = "PubMed"
42
+
43
+ def _run(self, history, **kwargs):
44
+ ans = self.get_input()
45
+ ans = " - ".join(ans["content"]) if "content" in ans else ""
46
+ if not ans:
47
+ return PubMed.be_output("")
48
+
49
+ Entrez.email = self._param.email
50
+ pubmedids = Entrez.read(Entrez.esearch(db='pubmed', retmax=self._param.top_n, term=ans))['IdList']
51
+ pubmedcnt = ET.fromstring(
52
+ Entrez.efetch(db='pubmed', id=",".join(pubmedids), retmode="xml").read().decode("utf-8"))
53
+ pubmed_res = [{"content": 'Title:' + child.find("MedlineCitation").find("Article").find(
54
+ "ArticleTitle").text + '\nUrl:<a href=" https://pubmed.ncbi.nlm.nih.gov/' + child.find(
55
+ "MedlineCitation").find("PMID").text + '">' + '</a>\n' + 'Abstract:' + child.find("MedlineCitation").find(
56
+ "Article").find("Abstract").find("AbstractText").text} for child in pubmedcnt.findall("PubmedArticle")]
57
+
58
+ if not pubmed_res:
59
+ return PubMed.be_output("")
60
+
61
+ df = pd.DataFrame(pubmed_res)
62
+ if DEBUG: print(df, ":::::::::::::::::::::::::::::::::")
63
+ return df
requirements.txt CHANGED
@@ -1,5 +1,6 @@
1
  Aspose.Slides==24.2.0
2
  BCEmbedding==0.1.3
 
3
  boto3==1.34.140
4
  botocore==1.34.140
5
  cachetools==5.3.3
 
1
  Aspose.Slides==24.2.0
2
  BCEmbedding==0.1.3
3
+ Bio==1.7.1
4
  boto3==1.34.140
5
  botocore==1.34.140
6
  cachetools==5.3.3
requirements_arm.txt CHANGED
@@ -151,3 +151,4 @@ duckduckgo_search==6.1.9
151
  google-generativeai==0.7.2
152
  groq==0.9.0
153
  wikipedia==1.4.0
 
 
151
  google-generativeai==0.7.2
152
  groq==0.9.0
153
  wikipedia==1.4.0
154
+ Bio==1.7.1
requirements_dev.txt CHANGED
@@ -136,3 +136,4 @@ duckduckgo_search==6.1.9
136
  google-generativeai==0.7.2
137
  groq==0.9.0
138
  wikipedia==1.4.0
 
 
136
  google-generativeai==0.7.2
137
  groq==0.9.0
138
  wikipedia==1.4.0
139
+ Bio==1.7.1