util.py

#!/usr/bin/python3.8

import dataclasses
import json
import functools
import os
import re
import string
import time

import anthropic
import groq
import mistralai.client
import mistralai.models

# infra - checkpoint while generating results

@dataclasses.dataclass
class CheckPointer:
  """Context manager for checkpointing while you accumulate results.
  """
  k_interval: int = 0
  time_interval: float = 60.0
  filestem: str = 'checkpt'
  last_chkpt_k: int = 0
  last_chkpt_time: float = 0.0
  k: int = 0
  results: list = dataclasses.field(default_factory=list)

  def __enter__(self):
    self.last_chkpt_time = time.time()
    print(f'initialize Checkpointing: {self.k=}')
    return self

  def __exit__(self, _type, _value, _traceback):
    self.flush(forced=True)
  
  def _num_ready_to_flush(self, forced):
    """Return (n, flush_now) where n is number of unflushed examples,
    and flush_now is true if that's enough to flush.
    """
    num_unflushed = self.k - self.last_chkpt_k
    if self.k_interval:
      if self.k >= self.last_chkpt_k + self.k_interval:
        return num_unflushed, True
    if self.time_interval:
      if time.time() >= self.last_chkpt_time + self.time_interval:
        return num_unflushed, True
    return num_unflushed, forced

  def flush(self, forced=False):
    """Write out unflushed results to a file.
    """
    num, flush_now = self._num_ready_to_flush(forced)
    if not flush_now or num==0:
      return
    elapsed = time.time() - self.last_chkpt_time
    hi = self.k
    lo = self.k - num
    file = f'{self.filestem}.{lo:04d}-{hi:04d}.jsonl'
    with open(file, 'w') as fp:
      for r in self.results[lo:hi]:
        fp.write(json.dumps(r) + '\n')
    print(f'write {file} with {num} outputs in {elapsed:.2f} sec')
    self.last_chkpt_k = self.k
    self.last_chkpt_time = time.time()
  
  def append(self, result):
    """Append a new result to the running list."""
    self.results.append(result)
    self.k += 1
    self.flush(forced=False)

# an LLM model

def llm(prompt, service='groq', model=None):
  if service == 'groq':
    if model is None: model='mixtral-8x7b-32768'
    client = groq.Groq(api_key=os.environ.get('GROQ_API_KEY'))
    completion = client.chat.completions.create(
      messages = [dict(role='user', content=prompt)],
      model=model)
    return completion.choices[0].message.content
  elif service == 'mistral':
    if model is None: model='open-mixtral-8x7b'
    client = mistralai.client.MistralClient(
      api_key=os.environ.get('MISTRAL_API_KEY'))
    chat_response = client.chat(
      messages=[mistralai.models.chat_completion.ChatMessage(
        role='user', content=prompt)],
      model=model)
    return chat_response.choices[0].message.content
  elif service == 'anthropic':
    if model is None: model='claude-3-haiku-20240307'
    client = anthropic.Anthropic(
      api_key=os.environ.get('ANTHROPIC_API_KEY'))
    message = client.messages.create(
      max_tokens=1024,
      model=model,
      messages=[{"role": "user", "content": prompt}])
    return message.content[0].text
  else:
    raise ValueError(f'invalid service {service}')

# run evaluation with a model

def run_eval(data, prompt_fn, parse_fn, delay=0.0, service='groq', **kw):
  with CheckPointer(**kw) as cp:
    for d in data:
      long_answer = llm(prompt_fn(d), service=service)
      short_answer = parse_fn(long_answer)
      d.update(ynu=short_answer, long_answer=long_answer)
      cp.append(d)
      time.sleep(delay)
  return cp.results

# normalize answers

# taken from https://github.com/allenai/bi-att-flow/blob/master/squad/evaluate-v1.1.py
def normalize_answer(s):
    """Lower text and remove punctuation, articles and extra whitespace."""
    def remove_articles(text):
        return re.sub(r'\b(a|an|the)\b', ' ', text)

    def white_space_fix(text):
        return ' '.join(text.split())

    def remove_punc(text):
        exclude = set(string.punctuation)
        return ''.join(ch for ch in text if ch not in exclude)

    def lower(text):
        return text.lower()

    return white_space_fix(remove_articles(remove_punc(lower(s))))