(entity_id, entity_name, pre_relations, pre_head, question, args)
| 107 | |
| 108 | |
| 109 | def relation_search_prune(entity_id, entity_name, pre_relations, pre_head, question, args): |
| 110 | sparql_relations_extract_head = sparql_head_relations % (entity_id) |
| 111 | head_relations = execurte_sparql(sparql_relations_extract_head) |
| 112 | head_relations = replace_relation_prefix(head_relations) |
| 113 | |
| 114 | sparql_relations_extract_tail= sparql_tail_relations % (entity_id) |
| 115 | tail_relations = execurte_sparql(sparql_relations_extract_tail) |
| 116 | tail_relations = replace_relation_prefix(tail_relations) |
| 117 | |
| 118 | if args.remove_unnecessary_rel: |
| 119 | head_relations = [relation for relation in head_relations if not abandon_rels(relation)] |
| 120 | tail_relations = [relation for relation in tail_relations if not abandon_rels(relation)] |
| 121 | |
| 122 | if pre_head: |
| 123 | tail_relations = list(set(tail_relations) - set(pre_relations)) |
| 124 | else: |
| 125 | head_relations = list(set(head_relations) - set(pre_relations)) |
| 126 | |
| 127 | head_relations = list(set(head_relations)) |
| 128 | tail_relations = list(set(tail_relations)) |
| 129 | total_relations = head_relations+tail_relations |
| 130 | total_relations.sort() # make sure the order in prompt is always equal |
| 131 | |
| 132 | if args.prune_tools == "llm": |
| 133 | prompt = construct_relation_prune_prompt(question, entity_name, total_relations, args) |
| 134 | |
| 135 | result = run_llm(prompt, args.temperature_exploration, args.max_length, args.opeani_api_keys, args.LLM_type) |
| 136 | flag, retrieve_relations_with_scores = clean_relations(result, entity_id, head_relations) |
| 137 | |
| 138 | elif args.prune_tools == "bm25": |
| 139 | topn_relations, topn_scores = compute_bm25_similarity(question, total_relations, args.width) |
| 140 | flag, retrieve_relations_with_scores = clean_relations_bm25_sent(topn_relations, topn_scores, entity_id, head_relations) |
| 141 | else: |
| 142 | model = SentenceTransformer('sentence-transformers/msmarco-distilbert-base-tas-b') |
| 143 | topn_relations, topn_scores = retrieve_top_docs(question, total_relations, model, args.width) |
| 144 | flag, retrieve_relations_with_scores = clean_relations_bm25_sent(topn_relations, topn_scores, entity_id, head_relations) |
| 145 | |
| 146 | if flag: |
| 147 | return retrieve_relations_with_scores |
| 148 | else: |
| 149 | return [] # format error or too small max_length |
| 150 | |
| 151 | |
| 152 | def entity_search(entity, relation, head=True): |
nothing calls this directly
no test coverage detected