trl-lib/trl-text-environment
9
1import os2import re3import copy4import time5 6import gradio as gr7from text_generation import Client8from transformers import load_tool9from share_btn import community_icon_html, loading_icon_html, share_js, share_btn_css10 11 12HF_TOKEN = os.environ.get("HF_TOKEN", None)13os.environ["HF_ALLOW_CODE_EVAL"] = "1"14print(HF_TOKEN)15 16FIM_PREFIX = "<fim_prefix>"17FIM_MIDDLE = "<fim_middle>"18FIM_SUFFIX = "<fim_suffix>"19 20FIM_INDICATOR = "<FILL_HERE>"21 22theme = gr.themes.Monochrome(23 primary_hue="indigo",24 secondary_hue="blue",25 neutral_hue="slate",26 radius_size=gr.themes.sizes.radius_sm,27 font=[28 gr.themes.GoogleFont("Open Sans"),29 "ui-sans-serif",30 "system-ui",31 "sans-serif",32 ],33)34 35tool = load_tool("vwxyzjn/pyserini-wikipedia-kilt-doc")36tool_fn = lambda x: tool(x).split("\n")[1][:600] # limit the amount if token, system_prompts37 38clients = {39 "StarCoderBase TriviaQA": [40 Client(41 "https://api-inference.huggingface.co/models/vwxyzjn/starcoderbase-triviaqa",42 headers={"Authorization": f"Bearer {HF_TOKEN}"},43 ),44 {"Wiki": tool_fn},45 """\46Answer the following question:47Q: In which branch of the arts is Patricia Neary famous?48A: Ballets49A2: <request><Wiki>Patricia Neary<call>Patricia Neary (born October 27, 1942) is an American ballerina, choreographer and ballet director, who has been particularly active in Switzerland. She has also been a highly successful ambassador for the Balanchine Trust, bringing George Balanchine's ballets to 60 cities around the globe.<response>50Result=Ballets<submit>51Q: Who won Super Bowl XX?52A: Chicago Bears53A2: <request><Wiki>Super Bowl XX<call>Super Bowl XX was an American football game between the National Football Conference (NFC) champion Chicago Bears and the American Football Conference (AFC) champion New England Patriots to decide the National Football League (NFL) champion for the 1985 season. The Bears defeated the Patriots by the score of 46–10, capturing their first NFL championship (and Chicago's first overall sports victory) since 1963, three years prior to the birth of the Super Bowl. Super Bowl XX was played on January 26, 1986 at the Louisiana Superdome in New Orleans.<response>54Result=Chicago Bears<submit>55""",56 ["Q: In which country is Oberhofen situated?", "Q: Irish Olympic champion Michelle smith was suspended in 1999 over drug allegations in which sport?"]57 ],58 "StarCoderBase GSM8K": [59 Client(60 "https://api-inference.huggingface.co/models/lvwerra/starcoderbase-gsm8k",61 headers={"Authorization": f"Bearer {HF_TOKEN}"},62 ),63 {"PythonInterpreter": load_tool("lvwerra/python-interpreter")},64 """\65Example of using a Python API to solve math questions. 66 67Q: Olivia has $23. She bought five bagels for $3 each. How much money does she have left?68 69<request><PythonInterpreter>70def solution():71 money_initial = 2372 bagels = 573 bagel_cost = 374 money_spent = bagels * bagel_cost75 money_left = money_initial - money_spent76 result = money_left77 return result78print(solution())79<call>72<response>80 81Result = 72 <submit>82""",83 ["Q: Tim has $400, and he received $1021. How much does he have?"]84 ],85}86 87def parse_tool_call(text, request_token="<request>", call_token="<call>"):88 """89 Parse request string. Expected format: <request><tool_name>query<call>90 """91 result = re.search(f"(?<={request_token}).*?(?={call_token})", text, re.DOTALL)92 93 # if we can't find a <request>/<call> span we return none94 if result is None:95 return None, None96 else:97 extracted_text = result.group()98 99 result = re.search(r"<(.*?)>", extracted_text)100 101 # if we can't find a tool name we return none102 if result is None:103 return None, None104 else:105 tool = result.group(1)106 107 # split off the tool name108 query = ">".join(extracted_text.split(">")[1:])109 110 return tool, query111 112 113 114def generate(115 prompt, system_prompt, version, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0,116):117 client, tools, _, _ = clients[version]118 temperature = float(temperature)119 if temperature < 1e-2:120 temperature = 1e-2121 top_p = float(top_p)122 fim_mode = False123 124 # TextEnv tool125 generate_kwargs = dict(126 temperature=temperature,127 max_new_tokens=max_new_tokens,128 top_p=top_p,129 repetition_penalty=repetition_penalty,130 do_sample=True,131 seed=42,132 stop_sequences=["<call>", "<submit>"]133 )134 generation_still_running = True135 request_idx = -1136 call_idx = -1137 response_idx = -1138 submit_idx = -1139 140 i = 0141 while generation_still_running:142 try:143 stream = client.generate_stream(system_prompt + prompt, **generate_kwargs)144 145 146 # call env phase147 output = system_prompt + prompt148 generation_start_idx = len(output)149 highlighted_output = [150 (prompt, "QUERY"),151 ]152 yield highlighted_output, output[generation_start_idx:]153 for response in stream:154 i += 1155 output += response.token.text156 tool, query = parse_tool_call(output[generation_start_idx:])157 158 if tool is not None and query is not None:159 # print("=====tool", i, tool, response, output)160 if tool not in tools:161 response = f"Unknown tool {tool}."162 try:163 response = tools[tool](query)164 output += response + "<response>"165 166 except Exception as error:167 response = f"Tool error: {str(error)}"168 169 if request_idx == -1:170 request_idx = output[generation_start_idx:].find("<request>")171 if call_idx == -1:172 call_idx = output[generation_start_idx:].find("<call>")173 if call_idx != -1:174 call_idx += len("<call>")175 if response_idx == -1:176 response_idx = output[generation_start_idx:].find("<response>")177 if response_idx != -1:178 response_idx += len("<response>")179 if submit_idx == -1:180 submit_idx = output[generation_start_idx:].find("<submit>")181 # I am sorry about the code182 print("-------", generation_start_idx, request_idx, call_idx, response_idx)183 highlighted_output = [184 (prompt, "QUERY"),185 (output[generation_start_idx:], "MODEL") if request_idx == -1 else ("", ""),186 (output[generation_start_idx:generation_start_idx+request_idx], "MODEL"),187 (output[generation_start_idx+request_idx:], "MODEL") if call_idx == -1 else ("", ""),188 (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "TOOL_REQUEST"),189 (output[generation_start_idx+call_idx:generation_start_idx+response_idx], "TOOL_CALL"),190 (output[generation_start_idx+response_idx:], "MODEL") if submit_idx != -1 else ("", ""),191 # (output[generation_start_idx:generation_start_idx+request_idx], ""),192 # (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "request"),193 # (output[generation_start_idx+call_idx:], "call"),194 ]195 print(i, highlighted_output, output[generation_start_idx:])196 yield highlighted_output, output[generation_start_idx:]197 198 # breakpoint()199 call_output = copy.deepcopy(output)200 print("start submit output")201 # response phase202 generate_kwargs["stop_sequences"] = ["<submit>"]203 stream = client.generate_stream(output, **generate_kwargs)204 for response in stream:205 output += response.token.text206 if submit_idx == -1:207 submit_idx = output[generation_start_idx:].find("<submit>")208 # print("-------", generation_start_idx, request_idx, call_idx, response_idx)209 highlighted_output = [210 (prompt, "QUERY"),211 (output[generation_start_idx:generation_start_idx+request_idx], "MODEL"),212 (output[generation_start_idx+request_idx:generation_start_idx+call_idx], "TOOL_REQUEST"),213 (output[generation_start_idx+call_idx:generation_start_idx+response_idx], "TOOL_CALL"),214 (output[generation_start_idx+response_idx:], "MODEL") if submit_idx != -1 else ("", ""),215 ]216 # print(highlighted_output, output[generation_start_idx:])217 yield highlighted_output, output[generation_start_idx:]218 print("-------", generation_start_idx, request_idx, call_idx, response_idx)219 print(highlighted_output, output[generation_start_idx:])220 221 return highlighted_output, output[generation_start_idx:]222 except Exception as e:223 if "loading" in str(e):224 gr.Warning("waiting for model to load... (this could take up to 20 minutes, after which things are much faster)")225 time.sleep(7)226 continue227 else:228 raise gr.Error(str(e)) 229 230 231examples = [232 "X_train, y_train, X_test, y_test = train_test_split(X, y, test_size=0.1)\n\n# Train a logistic regression model, predict the labels on the test set and compute the accuracy score",233 "// Returns every other value in the array as a new array.\nfunction everyOther(arr) {",234 "Poor English: She no went to the market. Corrected English:",235 "def alternating(list1, list2):\n results = []\n for i in range(min(len(list1), len(list2))):\n results.append(list1[i])\n results.append(list2[i])\n if len(list1) > len(list2):\n <FILL_HERE>\n else:\n results.extend(list2[i+1:])\n return results",236]237 238 239def process_example(args):240 for x in generate(args):241 pass242 return x243 244 245css = ".generating {visibility: hidden}"246 247monospace_css = """248#q-input textarea {249 font-family: monospace, 'Consolas', Courier, monospace;250}251"""252 253 254css += share_btn_css + monospace_css + ".gradio-container {color: black}"255 256 257description = """258<div style="text-align: center;">259 <img src="https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/images/textenv_demo_banner.png">260</div>261<div style="text-align: left;">262 <hr>263 <p>This is a demo to generate text the following StarCoderBase models fine-tuned using <a href="https://github.com/huggingface/trl/pull/424">TRL's TextEnvironment</a>:</p>264 <ul>265 <li><a href="https://huggingface.co/vwxyzjn/starcoderbase-triviaqa">StarCoderBase TriviaQA</a>: Uses a Wikipedia search index to answer trivia questions. It was trained on the TriviaQA dataset.</li>266 <li><a href="https://huggingface.co/lvwerra/starcoderbase-gsm8k">StarCoderBase GSM8K</a>: Uses a Python Interpreter to answer math questions. It was trained on the GSM8K dataset.</li>267 </ul>268</div>269"""270 271with gr.Blocks(theme=theme, analytics_enabled=False, css=css) as demo:272 with gr.Column():273 gr.Markdown(description)274 with gr.Row():275 version = gr.Dropdown(276 list(clients.keys()),277 value=list(clients.keys())[0],278 label="Model",279 info="Choose a model from the list",280 )281 282 with gr.Row():283 with gr.Column():284 instruction = gr.Textbox(285 value="Q: In which country is Oberhofen situated?",286 # placeholder="Enter your question here. E.g., Q: In which country is Oberhofen situated?",287 lines=2,288 label="Input",289 )290 submit = gr.Button("Generate", variant="primary")291 292 output = gr.HighlightedText(293 label="Output",294 color_map={"QUERY": "red", "TOOL_CALL": "green", "TOOL_RESPONSE": "blue", "MODEL": "pink"},295 )296 gr.Markdown("_Note:_ The trivia model is trained to give an answer first and then refine it with a Wiki call.")297 gr_examples = gr.Examples(298 examples=[example for client in clients.values() for example in client[3]],299 inputs=[instruction],300 cache_examples=False,301 )302 303 with gr.Row():304 with gr.Column():305 with gr.Accordion("Raw output", open=False):306 output2 = gr.Code(elem_id="q-output", lines=30, label="Raw output")307 with gr.Accordion("Advanced settings", open=False):308 with gr.Row():309 column_1, column_2 = gr.Column(), gr.Column()310 with column_1:311 temperature = gr.Slider(312 label="Temperature",313 value=0.2,314 minimum=0.0,315 maximum=1.0,316 step=0.05,317 interactive=True,318 info="Higher values produce more diverse outputs",319 )320 max_new_tokens = gr.Slider(321 label="Max new tokens",322 value=256,323 minimum=0,324 maximum=8192,325 step=64,326 interactive=True,327 info="The maximum numbers of new tokens",328 )329 with column_2:330 top_p = gr.Slider(331 label="Top-p (nucleus sampling)",332 value=0.90,333 minimum=0.0,334 maximum=1,335 step=0.05,336 interactive=True,337 info="Higher values sample more low-probability tokens",338 )339 repetition_penalty = gr.Slider(340 label="Repetition penalty",341 value=1.2,342 minimum=1.0,343 maximum=2.0,344 step=0.05,345 interactive=True,346 info="Penalize repeated tokens",347 )348 with gr.Accordion("Prompt", open=False):349 system_prompt = gr.Textbox(350 value=clients[list(clients.keys())[0]][2],351 label="System prompt",352 )353 version.select(354 lambda x: (clients[x][2]),355 inputs=[version],356 outputs=[system_prompt],357 )358 359 360 361 submit.click(362 generate,363 inputs=[instruction, system_prompt, version, temperature, max_new_tokens, top_p, repetition_penalty],364 outputs=[output, output2],365 )366demo.queue(concurrency_count=16).launch(debug=True)367