| 60 | return conv.get_prompt() |
| 61 | |
| 62 | def llm_llama(prompt: str) -> str: |
| 63 | CONTROLLER_ADDR = os.environ['CONTROLLER_ADDR'].split(',') |
| 64 | data = { |
| 65 | "inputs": prompt, |
| 66 | "parameters": { |
| 67 | "max_new_tokens": 256, |
| 68 | "do_sample": True, |
| 69 | 'temperature': 0.5, |
| 70 | 'frequency_penalty': 0, |
| 71 | 'presence_penalty': 0, |
| 72 | 'truncate': 4000, |
| 73 | } |
| 74 | } |
| 75 | if True or os.getenv('GREEDY'): |
| 76 | data['parameters']['do_sample'] = False |
| 77 | data['parameters'].pop('temperature') |
| 78 | print('greedy mode enabled') |
| 79 | for _ in range(5): |
| 80 | try: |
| 81 | response = requests.post( |
| 82 | random.choice(CONTROLLER_ADDR) + "/generate", |
| 83 | json=data, |
| 84 | timeout=120, |
| 85 | proxies={'http': '', 'https': ''}, |
| 86 | ) |
| 87 | print(response.content) |
| 88 | text = response.json()["generated_text"] |
| 89 | print(text) |
| 90 | return text.split('[INST]')[0].split('<|end_of_turn|>')[0].strip() |
| 91 | # if timeout or connection error, retry |
| 92 | except Timeout: |
| 93 | print("Timeout, retrying...") |
| 94 | except ConnectionError: |
| 95 | print("Connection error, retrying...") |
| 96 | except Exception: |
| 97 | traceback.print_exc() |
| 98 | try: |
| 99 | print(response) |
| 100 | print(response.text) |
| 101 | except: |
| 102 | pass |
| 103 | time.sleep(5) |
| 104 | else: |
| 105 | raise Exception("Timeout after 5 retries.") |
| 106 | |
| 107 | # Refresh traj log |
| 108 | def refresh(label: str): |