| 115 | raise Exception("Timeout after 3 retries.") |
| 116 | |
| 117 | def llm_tgi(prompt: str) -> str: |
| 118 | data = { |
| 119 | "inputs": prompt, |
| 120 | "parameters": { |
| 121 | "max_new_tokens": 256, |
| 122 | "do_sample": False, |
| 123 | 'truncate': 4000, |
| 124 | } |
| 125 | } |
| 126 | for _ in range(3): |
| 127 | try: |
| 128 | url = random.choice(CONTROLLER_ADDR) + "/generate" |
| 129 | print(f'Sending request to {url} ...') |
| 130 | response = requests.post( |
| 131 | url, |
| 132 | json=data, |
| 133 | timeout=120, |
| 134 | ) |
| 135 | text = response.json()["generated_text"] |
| 136 | print(text) |
| 137 | return text.split('[INST]')[0].split('<|end_of_turn|>')[0].strip() |
| 138 | # if timeout or connection error, retry |
| 139 | except Timeout: |
| 140 | print("Timeout, retrying...") |
| 141 | except ConnectionError: |
| 142 | print("Connection error, retrying...") |
| 143 | except Exception: |
| 144 | traceback.print_exc() |
| 145 | try: |
| 146 | print(response) |
| 147 | print(response.text) |
| 148 | except: |
| 149 | pass |
| 150 | time.sleep(5) |
| 151 | else: |
| 152 | raise Exception("Timeout after 3 retries.") |
| 153 | |
| 154 | def get_file_name(args, task_num): |
| 155 | if (len(args["output_path"]) > 0): |