-
Notifications
You must be signed in to change notification settings - Fork 0
/
config.py
75 lines (62 loc) · 1.77 KB
/
config.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
import typing as T
import os
# User libs
from pipeline.source import wikipedia, ollama, sparqlwrapper
from pipeline.metric import words, entities, class_count, entity_count, edges, distinct_edges, nodes, set_corpus
from pipeline.writer import write_text, write_dbpedia, write_ontology
# Schema definition for configuration file
ModelIdentifier = str
Query = str
Response = T.Any
SourceDictType = T.Dict[
ModelIdentifier,
T.Callable[[Query], Response]
]
TypeSpec = Response
SourceTypeType = T.Dict[TypeSpec, TypeSpec]
Metrics = T.Dict[
TypeSpec,
T.Dict[str, T.Callable[[Response], float]]
]
WriterTypes = T.Dict[TypeSpec, T.Callable[[Response, str], None]]
# == Configs
PROMPT: str = 'I am knowledgeable about pharmacy. Please provide all information about {}'
RESULT_DIR = 'results'
CORPUS_DIR = os.path.join(RESULT_DIR, 'corpus')
SOURCES: SourceDictType = {
'wikipedia': wikipedia,
'ollama-llama3': lambda query: ollama(query, model='llama3', prompt=PROMPT),
# 'ollama-phi3': lambda query: ollama(query, model='phi3', prompt=PROMPT),
# 'ollama-mistral': lambda query: ollama(query, model='mistral', prompt=PROMPT),
'dbpedia': sparqlwrapper,
}
SOURCE_TYPE: SourceTypeType = {
'wikipedia': 'text',
'ollama-llama3': 'text',
'ollama-phi3': 'text',
'ollama-mistral': 'text',
'dbpedia': 'dbpedia'
}
METRICS: Metrics = {
'text': {
'words': words,
'entities': entities
},
'dbpedia': {
'nodes': nodes,
'edges': edges,
'distinct_edges': distinct_edges
},
'ontology': {
'class_count': class_count,
'entity_count': entity_count
}
}
WRITERS: WriterTypes = {
'text': write_text,
'dbpedia': write_dbpedia,
'ontology': write_ontology
}
QUERIES = [
'Aspirin'
]