| import anthropic |
| from .common import TaskSpec, ParsedAnswer, Question |
| from .exceptions import GPTOutputParseException, GPTMaxTriesExceededException |
| import threading |
| from typing import List, Tuple, Union |
| from loguru import logger |
| from copy import deepcopy |
| import time |
| import os |
|
|
| class ClaudeModel(object): |
| def __init__(self, api_key:str, |
| task:TaskSpec, |
| model:str = "claude-3-haiku-20240307"): |
|
|
| self.claude_key:str = api_key |
| self.task:TaskSpec = task |
| self.model:str = model |
|
|
| def ask(self, payload:dict, n_choices=1) -> Tuple[List[dict], List[dict]]: |
| """ |
| args: |
| payload: json dictionary, prepared by `prepare_payload` |
| """ |
|
|
| def claude_thread(idx, payload, results): |
|
|
| |
| mod_payload = deepcopy(payload) |
|
|
| try: |
| raw_response = client.messages.create( |
| model="claude-3-haiku-20240307", |
| |
| messages = [mod_payload["messages"]], |
| max_tokens=mod_payload["max_tokens"], |
| ) |
| except Exception as e: |
| raise e |
|
|
| response = raw_response.dict() |
| response['content'] = response['content'][0]['text'] |
| message = {key: response[key] for key in ['role', 'content']} |
| metadata = response.copy() |
| del metadata["content"] |
| results[idx] = {"message": message, "metadata": metadata} |
| return |
|
|
| client = anthropic.Anthropic(api_key = self.claude_key) |
|
|
| assert n_choices >= 1 |
| results = [None] * n_choices |
| if n_choices > 1: |
| claude_jobs = [threading.Thread(target=claude_thread, |
| args=(idx, payload, results)) |
| for idx in range(n_choices)] |
| for job in claude_jobs: |
| job.start() |
| for job in claude_jobs: |
| job.join() |
| else: |
| claude_thread(0, payload, results) |
| messages:List[dict] = [ res["message"] for res in results] |
| metadata:List[dict] = [ res["metadata"] for res in results] |
| return messages, metadata |
|
|
|
|
| @staticmethod |
| def prepare_payload(question:Question, |
| max_tokens=1000, |
| verbose:bool=False, |
| prepend:Union[dict, None]=None, |
| **kwargs |
| ) -> dict: |
|
|
| content = [] |
| dic_list = question.get_json() |
| for dic in dic_list: |
| |
| if dic['type'] == 'text': |
| content.append(dic) |
|
|
| |
| elif dic['type'] == 'image_url': |
| base64enc_image = dic['image_url']['url'].split(',')[1] |
| if base64enc_image.startswith("/9j/"): |
| image_format = 'jpeg' |
| elif base64enc_image.startswith("iVBORw0KGgo"): |
| image_format = "png" |
| elif base64enc_image.startswith("R0lGOD") or base64enc_image.startswith("R0lGOD"): |
| image_format = "gif" |
| elif base64enc_image.startswith("UklGR"): |
| image_format = "webp" |
| else: |
| raise ValueError("Unknown format") |
|
|
| modified_dic = { |
| 'type' : "image", |
| 'source' : { |
| 'type' : "base64", |
| 'media_type' : f"image/{image_format}", |
| 'data' : base64enc_image |
| } |
| } |
|
|
| content.append(modified_dic) |
|
|
| payload = { |
| "messages": { |
| 'role': 'user', |
| 'content': content |
| }, |
| "max_tokens": max_tokens, |
| } |
|
|
|
|
| return payload |
|
|
|
|
| def rough_guess(self, question:Question, max_tokens=1000, |
| max_tries=1, query_id:int=0, |
| verbose=False, |
| **kwargs): |
| |
| p = self.prepare_payload(question, max_tokens = max_tokens, verbose=verbose, prepend=None, |
| model=self.model) |
|
|
| ok = False |
| while not ok: |
| response, meta_data = self.ask(p) |
| response = response [0] |
| try: |
| parsed_response = self.task.answer_type.parser(response["content"]) |
| except GPTOutputParseException as e: |
| if not os.path.exists('errors/'): |
| |
| os.makedirs('errors/') |
| error_saved = f'errors/{time.strftime("%Y-%m-%d-%H-%M-%S")}.json' |
| with open(error_saved, "w") as f: |
| f.write(p_ans.code) |
| |
| |
| reattempt += 1 |
| if reattempt > max_tries: |
| logger.error(f"max tries ({max_tries}) exceeded.") |
| raise GPTMaxTriesExceededException |
| |
| logger.warning(f"Reattempt #{reattempt} querying LLM") |
| continue |
| ok = True |
|
|
| return parsed_response, response, meta_data, p |
|
|
|
|
| def many_rough_guesses(self, num_threads:int, |
| question:Question, max_tokens=1000, |
| verbose=False, max_tries=1, |
| ) -> List[Tuple[ParsedAnswer, str, dict, dict]]: |
| """ |
| Args: |
| num_threads : number of independent threads. |
| all other arguments are same as those of `rough_guess()` |
| |
| Returns |
| List of elements, each element is a tuple following the |
| return signature of `rough_guess()` |
| """ |
|
|
| p = self.prepare_payload(question, max_tokens = max_tokens, verbose=verbose, prepend=None, |
| model=self.model) |
|
|
| n_choices = num_threads |
|
|
| ok = False |
| reattempt = 0 |
| while not ok: |
| response, meta_data = self.ask(p, n_choices=n_choices) |
| try: |
| parsed_response = [self.task.answer_type.parser(r["content"]) for r in response] |
| except GPTOutputParseException as e: |
| |
|
|
| reattempt += 1 |
| if reattempt > max_tries: |
| logger.error(f"max tries ({max_tries}) exceeded.") |
| raise GPTMaxTriesExceededException |
| |
| logger.warning(f"Reattempt #{reattempt} querying LLM") |
| continue |
| ok = True |
|
|
| return parsed_response, response, meta_data, p |
|
|
| def run_once(self, question:Question, max_tokens=1000, **kwargs): |
| q = self.task.first_question(question) |
| p_ans, ans, meta, p = self.rough_guess(q, max_tokens=max_tokens, **kwargs) |
| return p_ans, ans, meta, p |
