jojortz/llm4research-query-visualization
0
1import pprint2 3import pandas as pd4from uniflow.flow.client import TransformClient5from uniflow.flow.config import TransformOpenAIConfig6from uniflow.op.prompt import Context7 8from helpers import compare_strings_ignore_non_string9from visualize_upload import visualize10 11DEBUG = False12 13 14def cluster(query, answers_data):15 answers = []16 for answer in answers_data:17 answers.extend(answer["answer"])18 19 data = [Context(context=query, excerpts=answers)]20 21 instruction = """22# Task: I am a researcher with a query about research papers. I have a list of excerpts from those papers. I need you to cluster each of these excerpts into a category based on the query.23## Input:241. context: A brief query/context252. excerpts: An list of excerpts from research papers.26## Evaluation Steps:27### Step 128Go through each excerpt. For each excerpt, if there is an answer to the context/query that's not already captured by a category, create a category and add it to your category list. If the context has the word 'specific', make the category as specific as the excerpt. Repeat this process for each excerpt. The categories should be mutually exclusive.29### Step 230Once you've gone through all the excerpts and you have a list of categories, go through the excerpts a second time, and this time assign each excerpt to a category. A single excerpt can be assigned to multiple categories. If there is no information relevant to any of the categories, please categorize the excerpt as "None".31## Response Format: Your response should only include two fields below:321. categories: A list of all the generated categories. This is the output of Step 1 above.332. clusters: An object, with each category as a key, and a list of all the excerpts as strings that fall into that category as the value. This is the output of Step 2 above.34"""35 36 few_shot_examples = [37 # Context(38 # context="Which types of batteries are discussed?",39 # excerpts=[40 # "This investigation will shed lights on the tuneable chemical environments of transition-metal oxides for advanced cathode materials and promote the development of sodium-ion batteries.",41 # "Bi2Se3 was studied as a novel sodium-ion battery anode material because of its high theoretical capacity and high intrinsic conductivity.",42 # "Magnesium-ion batteries (MIBs) are considered strong candidates for next-generation energy-storage systems owing to their high theoretical capacity, divalent nature and the natural abundancy of magnesium (Mg) resources on Earth.",43 # "Magnesium-ion batteries (MIBs) have great potential in large-scale energy storage field with high capacity, excellent safety, and low cost.",44 # ],45 # categories=["Sodium-ion battery", "Magnesium-ion batteries"],46 # clusters={47 # "Sodium-ion battery": [48 # "This investigation will shed lights on the tuneable chemical environments of transition-metal oxides for advanced cathode materials and promote the development of sodium-ion batteries.",49 # "Bi2Se3 was studied as a novel sodium-ion battery anode material because of its high theoretical capacity and high intrinsic conductivity.",50 # ],51 # "Magnesium-ion batteries": [52 # "Magnesium-ion batteries (MIBs) are considered strong candidates for next-generation energy-storage systems owing to their high theoretical capacity, divalent nature and the natural abundancy of magnesium (Mg) resources on Earth.",53 # "Magnesium-ion batteries (MIBs) have great potential in large-scale energy storage field with high capacity, excellent safety, and low cost.",54 # ],55 # },56 # ),57 # Context(58 # context="Which 3D printing materials are discussed?",59 # excerpts=[60 # "The current state of materials development, including metal alloys, polymer composites, ceramics and concrete, was presented",61 # "To this end, this work designs a novel 3D printing phase change aggregate to prepare concrete with prominent thermal capacity and ductility.",62 # "In this study, 15 commercial pure titanium samples are processed under different conditions, and the 3D pore structures are characterized by X-ray tomography",63 # "In this study, a support-less ceramic printing (SLCP) process using a hydrogel bath was developed to facilitate the manufacture of complex bone substitutes.",64 # ],65 # categories=[66 # "metals",67 # "polymer composites",68 # "ceramics",69 # "concrete",70 # "phase change aggregate",71 # ],72 # clusters={73 # "metals": [74 # "The current state of materials development, including metal alloys, polymer composites, ceramics and concrete, was presented",75 # "In this study, 15 commercial pure titanium samples are processed under different conditions, and the 3D pore structures are characterized by X-ray tomography",76 # ],77 # "polymer composites": [78 # "The current state of materials development, including metal alloys, polymer composites, ceramics and concrete, was presented"79 # ],80 # "ceramics": [81 # "The current state of materials development, including metal alloys, polymer composites, ceramics and concrete, was presented",82 # "In this study, a support-less ceramic printing (SLCP) process using a hydrogel bath was developed to facilitate the manufacture of complex bone substitutes.",83 # ],84 # "concrete": [85 # "The current state of materials development, including metal alloys, polymer composites, ceramics and concrete, was presented",86 # "To this end, this work designs a novel 3D printing phase change aggregate to prepare concrete with prominent thermal capacity and ductility.",87 # ],88 # "phase change aggregate": [89 # "To this end, this work designs a novel 3D printing phase change aggregate to prepare concrete with prominent thermal capacity and ductility."90 # ],91 # },92 # ),93 ]94 95 num_thread_batch_size = 196 97 config = TransformOpenAIConfig()98 config.prompt_template.instruction = instruction99 config.prompt_template.few_shot_prompt = few_shot_examples100 config.model_config.model_name = "gpt-4-1106-preview"101 config.model_config.response_format = {"type": "json_object"}102 config.model_config.num_call = 1103 config.model_config.temperature = 0.0104 config.model_config.num_thread = num_thread_batch_size105 config.model_config.batch_size = num_thread_batch_size106 107 cluster_client = TransformClient(config)108 109 output = cluster_client.run(data)110 if DEBUG:111 pprint.pprint(output)112 output_clusters = answers_data113 clusters = output[0]["output"][0]["response"][0]["clusters"]114 output_answer_category = []115 116 for idx, paper in enumerate(answers_data):117 # Initialize an empty list to store the categories for each answer118 categories_per_answer = []119 120 # Iterate over each answer121 for ans in paper["answer"]:122 categories = []123 # Iterate over each category in clusters124 for category, texts in clusters.items():125 # Check if the answer is in any of the texts related to the category126 if any(compare_strings_ignore_non_string(ans, text) for text in texts):127 if category not in categories_per_answer:128 categories.append(category)129 output_answer_category.append(130 {"paper": paper["paper"], "answer": ans, "category": category}131 )132 if len(categories) == 0:133 categories.append("None")134 categories_per_answer.extend(categories)135 136 output_clusters[idx]["categories"] = categories_per_answer137 for output_cluster in output_clusters:138 if len(output_cluster["categories"]) == 0:139 output_cluster["categories"].append("None")140 df = create_category_df(output_clusters, answers_data)141 output_answer_category_df = pd.DataFrame(output_answer_category)142 visualize_output = visualize(output_clusters)143 144 return [output_clusters, df, visualize_output, output_answer_category_df]145 146 147def create_category_df(cluster_output, answers_data):148 pd_data = {149 "Paper": [],150 "Excerpts": [],151 "Categories": [],152 }153 for i, paper in enumerate(cluster_output):154 pd_data["Paper"].append(paper["paper"])155 pd_data["Excerpts"].append(", ".join(answers_data[i]["answer"]))156 pd_data["Categories"].append(", ".join(paper["categories"]))157 158 df = pd.DataFrame(pd_data)159 return df160