feat: add audio narration (updated)

alembic/versions/c176288cb508_add_audio_info.py

@@ -0,0 +1,47 @@
+"""Add audio info.
+
+Revision ID: c176288cb508
+Revises: 8713b142f5de
+Create Date: 2023-08-31 00:25:04.889325
+
+"""
+import sqlalchemy as sa
+
+from alembic import op
+from openadapt.models import ForceFloat
+
+# revision identifiers, used by Alembic.
+revision = "c176288cb508"
+down_revision = "8713b142f5de"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+ # ### commands auto generated by Alembic - please adjust! ###
+ op.create_table(
+ "audio_info",
+ sa.Column("id", sa.Integer(), nullable=False),
+ sa.Column("flac_data", sa.LargeBinary(), nullable=True),
+ sa.Column("transcribed_text", sa.String(), nullable=True),
+ sa.Column(
+ "recording_timestamp",
+ ForceFloat(precision=10, scale=2, asdecimal=False),
+ nullable=True,
+ ),
+ sa.Column("sample_rate", sa.Integer(), nullable=True),
+ sa.Column("words_with_timestamps", sa.Text(), nullable=True),
+ sa.ForeignKeyConstraint(
+ ["recording_timestamp"],
+ ["recording.timestamp"],
+ name=op.f("fk_audio_info_recording_timestamp_recording"),
+ ),
+ sa.PrimaryKeyConstraint("id", name=op.f("pk_audio_info")),
+ )
+ # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+ # ### commands auto generated by Alembic - please adjust! ###
+ op.drop_table("audio_info")
+ # ### end Alembic commands ###

openadapt/db/crud.py

@@ -4,6 +4,7 @@
 """
 
 from typing import Any
+import json
 
 from loguru import logger
 import sqlalchemy as sa
@@ -12,6 +13,7 @@
 from openadapt.db.db import BaseModel, Session
 from openadapt.models import (
  ActionEvent,
+ AudioInfo,
  MemoryStat,
  PerformanceStat,
  Recording,
@@ -411,3 +413,22 @@ def new_session() -> None:
  if db:
  db.close()
  db = Session()
+
+
+def insert_audio_info(
+ audio_data: bytes,
+ transcribed_text: str,
+ recording_timestamp: float,
+ sample_rate: int,
+ word_list: list,
+) -> None:
+ """Create an AudioInfo entry in the database."""
+ audio_info = AudioInfo(
+ flac_data=audio_data,
+ transcribed_text=transcribed_text,
+ recording_timestamp=recording_timestamp,
+ sample_rate=sample_rate,
+ words_with_timestamps=json.dumps(word_list),
+ )
+ db.add(audio_info)
+ db.commit()

openadapt/db/db.py

@@ -5,7 +5,7 @@
 
 from dictalchemy import DictableModel
 from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import scoped_session, sessionmaker
 from sqlalchemy.schema import MetaData
 import sqlalchemy as sa
 
@@ -67,4 +67,5 @@ def get_base(engine: sa.engine) -> sa.engine:
 
 engine = get_engine()
 Base = get_base(engine)
-Session = sessionmaker(bind=engine)
+session_factory = sessionmaker(bind=engine)
+Session = scoped_session(session_factory)

openadapt/models.py

@@ -61,6 +61,8 @@ class Recording(db.Base):
  order_by="WindowEvent.timestamp",
  )
 
+ audio_info = sa.orm.relationship("AudioInfo", back_populates="recording")
+
  _processed_action_events = None
 
  @property
@@ -378,6 +380,21 @@ def get_active_window_event(cls: "WindowEvent") -> "WindowEvent":
  return WindowEvent(**window.get_active_window_data())
 
 
+class AudioInfo(db.Base):
+ """Class representing the audio from a recording in the database."""
+
+ __tablename__ = "audio_info"
+
+ id = sa.Column(sa.Integer, primary_key=True)
+ flac_data = sa.Column(sa.LargeBinary)
+ transcribed_text = sa.Column(sa.String)
+ recording_timestamp = sa.Column(sa.ForeignKey("recording.timestamp"))
+ sample_rate = sa.Column(sa.Integer)
+ words_with_timestamps = sa.Column(sa.Text)
+
+ recording = sa.orm.relationship("Recording", back_populates="audio_info")
+
+
 class PerformanceStat(db.Base):
  """Class representing a performance statistic in the database."""
 

openadapt/record.py

@@ -4,11 +4,16 @@
 
  $ python openadapt/record.py "<description of task to be recorded>"
 
+To record audio:
+
+ $ python openadapt/record.py "<description of task to be recorded>" --enable_audio
+
 """
 
 from collections import namedtuple
 from functools import partial, wraps
 from typing import Any, Callable, Union
+import io
 import multiprocessing
 import os
 import queue
@@ -24,7 +29,11 @@
 from tqdm import tqdm
 import fire
 import mss.tools
+import numpy as np
 import psutil
+import sounddevice
+import soundfile
+import whisper
 
 from openadapt import config, utils, window
 from openadapt.db import crud
@@ -804,15 +813,101 @@ def read_mouse_events(
  mouse_listener.stop()
 
 
+def record_audio(
+ terminate_event: multiprocessing.Event,
+ recording_timestamp: float,
+) -> None:
+ """Record audio narration during the recording and store data in database.
+
+ Args:
+ terminate_event: The event to signal termination of event reading.
+ recording_timestamp: The timestamp of the recording.
+ """
+ utils.configure_logging(logger, LOG_LEVEL)
+ utils.set_start_time(recording_timestamp)
+
+ audio_frames = [] # to store audio frames
+
+ def audio_callback(
+ indata: np.ndarray, frames: int, time: Any, status: sounddevice.CallbackFlags
+ ) -> None:
+ """Callback function used when new audio frames are recorded.
+
+ Note: time is of type cffi.FFI.CData, but since we don't use this argument
+ and we also don't use the cffi library, the Any type annotation is used.
+ """
+ # called whenever there is new audio frames
+ audio_frames.append(indata.copy())
+
+ # open InputStream and start recording while ActionEvents are recorded
+ audio_stream = sounddevice.InputStream(
+ callback=audio_callback, samplerate=16000, channels=1
+ )
+ logger.info("Audio recording started.")
+ audio_stream.start()
+ terminate_event.wait()
+ audio_stream.stop()
+ audio_stream.close()
+
+ # Concatenate into one Numpy array
+ concatenated_audio = np.concatenate(audio_frames, axis=0)
+ # convert concatenated_audio to format expected by whisper
+ converted_audio = concatenated_audio.flatten().astype(np.float32)
+
+ # Convert audio to text using OpenAI's Whisper
+ logger.info("Transcribing audio...")
+ model = whisper.load_model("base")
+ result_info = model.transcribe(converted_audio, word_timestamps=True, fp16=False)
+ logger.info(f"The narrated text is: {result_info['text']}")
+ # empty word_list if the user didn't say anything
+ word_list = []
+ # segments could be empty
+ if len(result_info["segments"]) > 0:
+ # there won't be a 'words' list if the user didn't say anything
+ if "words" in result_info["segments"][0]:
+ word_list = result_info["segments"][0]["words"]
+
+ # compress and convert to bytes to save to database
+ logger.info(
+ "Size of uncompressed audio data: {} bytes".format(converted_audio.nbytes)
+ )
+ # Create an in-memory file-like object
+ file_obj = io.BytesIO()
+ # Write the audio data using lossless compression
+ soundfile.write(
+ file_obj, converted_audio, int(audio_stream.samplerate), format="FLAC"
+ )
+ # Get the compressed audio data as bytes
+ compressed_audio_bytes = file_obj.getvalue()
+
+ logger.info(
+ "Size of compressed audio data: {} bytes".format(len(compressed_audio_bytes))
+ )
+
+ file_obj.close()
+
+ # To decompress the audio and restore it to its original form:
+ # restored_audio, restored_samplerate = sf.read(
+ # io.BytesIO(compressed_audio_bytes))
+
+ # Create AudioInfo entry
+ crud.insert_audio_info(
+ compressed_audio_bytes,
+ result_info["text"],
+ recording_timestamp,
+ int(audio_stream.samplerate),
+ word_list,
+ )
+
+
 @logger.catch
 @trace(logger)
-def record(
- task_description: str,
-) -> None:
+def record(task_description: str, enable_audio: bool = False) -> None:
  """Record Screenshots/ActionEvents/WindowEvents.
 
  Args:
  task_description: A text description of the task to be recorded.
+ enable_audio: a flag to enable or disable audio recording (default: False)
  """
  logger.info(f"{task_description=}")
 
@@ -943,6 +1038,13 @@ def record(
  )
  mem_plotter.start()
 
+ if enable_audio:
+ audio_recorder = threading.Thread(
+ target=record_audio,
+ args=(terminate_event, recording_timestamp),
+ )
+ audio_recorder.start()
+
  # TODO: discard events until everything is ready
 
  collect_stats()
@@ -972,6 +1074,9 @@ def record(
  screen_event_writer.join()
  action_event_writer.join()
  window_event_writer.join()
+ if enable_audio:
+ audio_recorder.join()
+
  terminate_perf_event.set()
 
  if PLOT_PERFORMANCE:

openadapt/strategies/mixins/openai.py

@@ -11,7 +11,6 @@ class MyReplayStrategy(OpenAIReplayStrategyMixin):
 
 from loguru import logger
 import openai
-import tiktoken
 
 from openadapt import cache, config, models
 from openadapt.strategies.base import BaseReplayStrategy
@@ -29,7 +28,6 @@ class MyReplayStrategy(OpenAIReplayStrategyMixin):
 MODEL_NAME = "gpt-4"
 
 openai.api_key = config.OPENAI_API_KEY
-encoding = tiktoken.get_encoding("cl100k_base")
 
 
 class OpenAIReplayStrategyMixin(BaseReplayStrategy):
@@ -187,50 +185,3 @@ def _get_completion(prompt: str) -> str:
  logger.debug(f"appending assistant_message=\n{pformat(assistant_message)}")
  messages.append(assistant_message)
  return messages
-
-
-# XXX TODO not currently in use
-# https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
-def num_tokens_from_messages(messages: list, model: str = "gpt-3.5-turbo-0301") -> int:
- """Returns the number of tokens used by a list of messages."""
- try:
- encoding = tiktoken.encoding_for_model(model)
- except KeyError:
- logger.info("Warning: model not found. Using cl100k_base encoding.")
- encoding = tiktoken.get_encoding("cl100k_base")
- if model == "gpt-3.5-turbo":
- logger.info(
- "Warning: gpt-3.5-turbo may change over time. Returning num tokens "
- "assuming gpt-3.5-turbo-0301."
- )
- return num_tokens_from_messages(messages, model="gpt-3.5-turbo-0301")
- elif model == "gpt-4":
- logger.info(
- "Warning: gpt-4 may change over time. Returning num tokens "
- "assuming gpt-4-0314."
- )
- return num_tokens_from_messages(messages, model="gpt-4-0314")
- elif model == "gpt-3.5-turbo-0301":
- tokens_per_message = (
- 4 # every message follows <|start|>{role/name}\n{content}<|end|>\n
- )
- tokens_per_name = -1 # if there's a name, the role is omitted
- elif model == "gpt-4-0314":
- tokens_per_message = 3
- tokens_per_name = 1
- else:
- raise NotImplementedError(
- f"""num_tokens_from_messages() is not implemented for model "
- "{model}. See "
- "https://github.com/openai/openai-python/blob/main/chatml.md for "
- information on how messages are converted to tokens."""
- )
- num_tokens = 0
- for message in messages:
- num_tokens += tokens_per_message
- for key, value in message.items():
- num_tokens += len(encoding.encode(value))
- if key == "name":
- num_tokens += tokens_per_name
- num_tokens += 3 # every reply is primed with <|start|>assistant<|message|>
- return num_tokens

openadapt/visualize.py

@@ -12,7 +12,7 @@
 import click
 
 from openadapt import config
-from openadapt.db.crud import get_latest_recording, get_recording
+from openadapt.db.crud import get_audio_info, get_latest_recording, get_recording
 from openadapt.events import get_events
 from openadapt.utils import (
  EMPTY,
@@ -147,6 +147,10 @@ def main(timestamp: str, notify: bool = True) -> None:
  scrub.scrub_text(recording.task_description)
  logger.debug(f"{recording=}")
 
+ audio_info = row2dict(get_audio_info(recording))
+ # don't display the FLAC data
+ del audio_info["flac_data"]
+
  meta = {}
  action_events = get_events(recording, process=PROCESS_EVENTS, meta=meta)
  event_dicts = rows2dicts(action_events)