OpenWormLLM / datasources.py
pgleeson's picture
Update datasources.py
515ef92 verified
Raw
History Blame
937 Bytes
# Based on: https://github.com/allenai/s2-folks/tree/main/examples/python
#
#
import os
import re
FORMAT_TEXT = 'Text'
FORMAT_DICT = 'Dict'
FORMAT_JSON = 'JSON'
FORMATS = (FORMAT_TEXT, FORMAT_DICT)
DS_WORMNEUROATLAS = 'WormNeuroAtlas'
import wormneuroatlas as wa
import pprint
pp = pprint.PrettyPrinter(depth=6)
def query_data_source(text, source, format):
if source == DS_WORMNEUROATLAS:
atlas = wa.NeuroAtlas()
ds, ss = atlas.everything_about(text, return_values=True, return_text=True, print_text=False)
if format==FORMAT_TEXT:
info = '''Extracting info on neuron %s from %s...
'''%(text, source)
info += """```%s```"""%ss
elif format==FORMAT_DICT:
info = """```
%s```"""%pp.pformat(ds)
elif format==FORMAT_JSON:
import json
jj = json.dumps(ds)
info = """```
%s```"""%pp.pformat(jj)
return info