repo_description: AI-powered tools to enhance Anki flashcards with explanations, mnemonics, illustrations, and adaptive learning for medical school and beyond
repo_topics: anki dalle-2 dalle-3 illustrations learning llm medical mnemonics openai spaced-repetition stable-diffusion summary
code: """
  Medical mnemonics generation tool using LLMs.
  
  This module provides functionality to generate and manage mnemonic devices for medical concepts
  using large language models. It includes semantic search to find similar existing mnemonics
  and an interactive CLI interface for selecting from multiple generated options.
  """
  from pathlib import Path
  import litellm
  import json
  from typing import List
  import numpy as np
  import fire
  from sklearn.metrics.pairwise import cosine_similarity
  import asyncio
  import joblib
  from tqdm.asyncio import tqdm_asyncio
  
  from utils.llm import model_name_matcher
  
  from prompt_toolkit import PromptSession
  from prompt_toolkit.completion import WordCompleter
  from prompt_toolkit.application import Application
  from prompt_toolkit.key_binding import KeyBindings
  from prompt_toolkit.layout import Layout
  from prompt_toolkit.layout.containers import HSplit, Window
  from prompt_toolkit.layout.controls import FormattedTextControl
  
  
  from iterator_cacher import IteratorCacher
  
  # anchor_file = Path("./author_dir/anchors.json")
  anchor_file = Path("/home/g/Documents/Perso/Bordel/git_repository/AnkiAiUtils/author_dir/anchors.json")
  new_anchor_file = Path("/home/g/Documents/Perso/Bordel/git_repository/AnkiAiUtils/author_dir/new_anchors.json")
  
  system_prompt = """
  You are my best assistant. Your task is to generate new mnemonics to help me in medical school.
  
  You follow these rules:
  - Answer in the same language as the question.
  - Answer only 1 mnemonic.
  - The new mnemonics have to be easy to remember, so funny or tragic.
  - The new mnemonics can only contain common words and can't use technical words.
  - The new mnemonics can't contain digits.
  - Directly answer the mnemonics, don't try to be polite, don't acknowledge those rules. Don't wrap your answer. Don't wrap the mnemonic in quotation signs.
  - After your mnemonic, mention a single sentence between parentheses where you explain your reasonning. This will help me memorize it later.
  - Your full answer must remain in a single line.
  """.strip()
  
  def prompt(
      examples: dict,
      question: str,
      ) -> str:
      """
      Generate a prompt for the LLM to create a mnemonic device.
  
      Parameters
      ----------
      examples : dict
          Dictionary of existing mnemonics to use as examples, where keys are concepts
          and values are their corresponding mnemonic devices
      question : str
          The concept/topic for which to generate a new mnemonic device
  
      Returns
      -------
      str
          The formatted prompt containing the system instructions, examples, and question
      """
      message = f"La notion pour laquelle j'ai besoin d'un moyen mnémotechnique est: '{question}'"
      message += "\nVoici des exemples de moyen mnémotechniques que j'utilise déjà:\n'''\n"
      ex = []
      for k, v in examples.items():
          ex.append(f"- '{k}': '{v}'")
  
      # ex = ex[::-1]  # reverse sorting order to have the end be closer to the question
  
      message += "\n".join(ex)
  
      message += f"\n'''\n\nMaintenant, suit les règles que je t'ai donné et suggère moi un moyen mnemotechnique adéquats. Je rappelle que la notion est: '{question}'"
      return message
  
  anchors = json.loads(anchor_file.read_text())
  assert anchors
  try:
      new_anchors = json.loads(new_anchor_file.read_text())
  except Exception as err:
      print(f"Not loading new anchors: '{err}'")
      new_anchors = {}
  
  
  mem_object = joblib.Memory(".cache/iterator_cached_embeddings")
  def embed(
      text_list: List[str],
      model: str = "openai/text-embedding-3-small",
      ):
      """
      Generate embeddings for a list of text strings using a specified model.
  
      Parameters
      ----------
      text_list : List[str]
          List of text strings to embed
      model : str, optional
          Name of the embedding model to use, by default "openai/text-embedding-3-small"
  
      Returns
      -------
      numpy.ndarray
          Array of embeddings, one per input text
      """
      assert text_list
      assert all(t.strip() for t in text_list)
      vec = litellm.embedding(
              model=model,
              input=text_list,
              )
      embeds = [
          np.array(elem["embedding"]).reshape(1, -1)
          for elem in vec.to_dict()["data"]
      ]
      embeds = np.array(embeds)
      return embeds
  
  embed = IteratorCacher(
      memory_object=mem_object,
      iter_list=["text_list"],
      verbose=True,
      res_to_list = lambda ar: ar.tolist(),
  )(embed)
  
  
  def cli(
      top_k: int,
      n_gen: int,
      query: str,
      model: str,
      embed_model: str,
      ):
      """
      Run the interactive CLI interface for mnemonic generation.
  
      Parameters
      ----------
      top_k : int
          Number of similar existing mnemonics to use as examples
      n_gen : int
          Number of new mnemonics to generate per query
      query : str
          Initial query to process, if any
      model : str
          Name of the LLM model to use for generation
      embed_model : str
          Name of the model to use for text embeddings
      """
      session = PromptSession()
  
      if model.startswith("openrouter"):
          model_params = litellm.get_supported_openai_params(
              model_name_matcher(
                  model.split("/", 1)[1]
                  )
              )
      else:
          model_params = litellm.get_supported_openai_params(model)
  
      while True:
  
          # refresh db
          for k, v in new_anchors.items():
              assert k not in anchors.keys()
              assert v not in anchors.values()
          anchors_list = sorted(list(anchors.keys()) + list(new_anchors.keys()))
          db = np.array(
                  embed(
                  text_list=anchors_list,
                  model=embed_model,
              )
          )
  
          if query is not None:
              ans, query = query, None
              query = None
          else:
              ans = session.prompt(
                  "Enter the string your want mnemonics for:\n>",
                  vi_mode=True,
                  multiline=False,
                  completer=WordCompleter(anchors_list),
              ).strip()
          assert ans
          if ans in anchors_list:
              if ans in anchors:
                  print(f"Already part of anchors: '{anchors[ans]}'")
              elif ans in new_anchors:
                  print(f"Already part of anchors: '{new_anchors[ans]}'")
              else:
                  raise ValueError(ans)
  
          vec = embed(text_list=[ans])
          assert len(vec) == 1
          vec = vec[0]
          if isinstance(vec, list):
              vec = np.array(vec)
  
          sims = cosine_similarity(
              vec.squeeze().reshape(1, -1),
              db.squeeze(),
          )
  
          best_index = np.argsort(sims).tolist()[0][-top_k:]
          best_k = [anchors_list[b] for b in best_index]
          best_a = [anchors[k] if k in anchors else new_anchors[k] for k in best_k]
          best = {k:v for k, v in zip(best_k, best_a)}
  
  
          messages = [
              {
                  "role": "system",
                  "content": system_prompt,
              },
              {
                  "role": "user",
                  "content": prompt(examples=best, question=ans),
              }
          ]
  
          if "n" in model_params:
              resp = litellm.completion(
                  model=model,
                  messages=messages,
                  num_retries=3,
                  verbose=True,
                  n=n_gen,
                  temperature=1,
              ).json()
              texts = [elem["message"]["content"] for elem in resp["choices"]]
          else:
              async def async_completion(model, messages):
                  loop = asyncio.get_event_loop()
                  resp = await loop.run_in_executor(None, lambda: litellm.completion(
                      model=model,
                      messages=messages,
                      num_retries=3,
                      verbose=True,
                      temperature=1,
                  ).json())
                  return resp["choices"][0]["message"]["content"]
              async def generate_texts_async(n_gen, model, messages):
                  tasks = [async_completion(model, messages) for _ in range(n_gen)]
                  return await tqdm_asyncio.gather(*tasks, desc="Generating")
              def generate_texts(n_gen, model, messages):
                  """
                  Generate multiple mnemonic texts using async calls to the LLM.
  
                  Parameters
                  ----------
                  n_gen : int
                      Number of texts to generate
                  model : str
                      Name of the LLM model to use
                  messages : List[Dict]
                      Prompt messages to send to the model
  
                  Returns
                  -------
                  List[str]
                      List of generated mnemonic texts
                  """
                  return asyncio.run(generate_texts_async(n_gen, model, messages))
  # Usage (inside your synchronous function):
              texts = asyncio.run(generate_texts_async(n_gen, model, messages))
  
  
          texts.append("Cancel")
  
          selected_index = [0]
          def get_menu_text():
              """
              Format the menu text for the CLI interface.
  
              Returns
              -------
              List[Tuple[str, str]]
                  List of formatted text entries with their styles
              """
              out = []
              for i, item in enumerate(texts):
                  n = ' ' * len(str(n_gen)) if item == "Cancel" else str(i + 1)
                  sign = ">" if i == selected_index[0] else " "
                  line = f"{sign} {n} {item}\n\n"
                  out.append(("class:menu-item", line))
              return out
  
          kb = KeyBindings()
          @kb.add("k")
          @kb.add("up")
          def move_up(event):
              selected_index[0] = max(selected_index[0] - 1, 0)
  
          @kb.add("j")
          @kb.add("down")
          def move_down(event):
              selected_index[0] = min(selected_index[0] + 1, len(texts) - 1)
  
          @kb.add("enter")
          def select(event):
              event.app.exit()
  
          @kb.add("c-c")
          @kb.add("q")
          def quit(event):
              print("Quit.")
              raise SystemExit()
  
          for i in range(10):
              @kb.add(str(i))
              def _(event, i=i):
                  i = min(i, len(texts))
                  i = max(0, i)
                  selected_index[0] = i
  
  
          menu_control = FormattedTextControl(get_menu_text)
          menu_window = Window(
              content=menu_control,
              always_hide_cursor=True,
              wrap_lines=True,
          )
          root_container = HSplit([menu_window])
          layout = Layout(root_container)
          app = Application(layout=layout, key_bindings=kb, full_screen=False)
          print("\n")
          try:
              app.run()
          except EOFError:
              raise SystemExit()
  
          new_anc = texts[selected_index[0]]
          if new_anc == "Cancel":
              print("Cancelled")
              continue
  
          # store to new_anchors
          new_anchors[ans] = new_anc
          with open(new_anchor_file, "w") as f:
              json.dump(new_anchors, f, indent=4, ensure_ascii=False)
  
  def start(
      top_k: int = 100,
      n_gen: int = 10,
      query: str = None,
      model: str = "openrouter/anthropic/claude-3.5-sonnet:beta",
      embed_model: str = "openai/text-embedding-3-small",
      gui: bool = False,
      ):
      """
      Entry point for the mnemonic generation tool.
  
      Parameters
      ----------
      top_k : int, optional
          Number of similar existing mnemonics to use, by default 100
      n_gen : int, optional
          Number of new mnemonics to generate per query, by default 10
      query : str, optional
          Initial query to process, by default None
      model : str, optional
          LLM model name, by default "openrouter/anthropic/claude-3.5-sonnet:beta"
      embed_model : str, optional
          Embedding model name, by default "openai/text-embedding-3-small"
      gui : bool, optional
          Whether to use GUI interface (not implemented), by default False
      """
      if not gui:
          cli(
              top_k=top_k,
              n_gen=n_gen,
              query=query,
              model=model,
              embed_model=embed_model,
          )
      else:
          raise NotImplementedError()
  
  if __name__ == "__main__":
      fire.Fire(start)
  