demo_llama.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This quickstart tutorial will walk you through how to use an Extended Mind Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start with a simple question-answer example, and see if we can eliminate a simple model hallucination."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"normalcomputing/extended-mind-llama-2-7b-chat\")\n",
    "model = AutoModelForCausalLM.from_pretrained(\"normalcomputing/extended-mind-llama-2-7b-chat\", trust_remote_code=True).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"When did Alexander Grothendieck get his French citizenship?\"\n",
    "inputs = tokenizer(query, return_tensors=\"pt\").input_ids\n",
    "\n",
    "outputs_hf = model.generate(inputs.to(device), max_length= inputs.size(-1) + 50,)\n",
    "print(tokenizer.decode(outputs_hf[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Oops! Not quite right. Now, let's give the model the information it needs as memories. We can pass memories in the `from_pretrained()` method, or easily set them afterwards as below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ag_wiki_entry = \"\"\"Alexander Grothendieck (/ˈɡroʊtəndiːk/; German pronunciation: [ˌalɛˈksandɐ ˈɡʁoːtn̩ˌdiːk] (listen); French: [ɡʁɔtɛndik]; 28 March 1928 – 13 November 2014) was a stateless (and then, since 1971, French) mathematician who became the leading figure in the creation of modern algebraic geometry.[7][8] His research extended the scope of the field and added elements of commutative algebra, homological algebra, sheaf theory, and category theory to its foundations, while his so-called \"relative\" perspective led to revolutionary advances in many areas of pure mathematics.[7][9] He is considered by many to be the greatest mathematician of the twentieth century.[10][11]\n",
    "\n",
    "Grothendieck began his productive and public career as a mathematician in 1949. In 1958, he was appointed a research professor at the Institut des hautes études scientifiques (IHÉS) and remained there until 1970, when, driven by personal and political convictions, he left following a dispute over military funding. He received the Fields Medal in 1966 for advances in algebraic geometry, homological algebra, and K-theory.[12] He later became professor at the University of Montpellier[1] and, while still producing relevant mathematical work, he withdrew from the mathematical community and devoted himself to political and religious pursuits (first Buddhism and later, a more Christian vision).[13] In 1991, he moved to the French village of Lasserre in the Pyrenees, where he lived in seclusion, still working tirelessly on mathematics and his philosophical and religious thoughts until his death in 2014.[14]\n",
    "\"\"\"\n",
    "memories = tokenizer(ag_wiki_entry).input_ids\n",
    "model.memory_ids = memories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = model.generate(inputs.to(device), max_length=60, topk=2, output_retrieved_memory_idx=True, output_attentions=True, return_dict_in_generate=True,)\n",
    "print(tokenizer.decode(outputs['sequences'][0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nice! Let's check out which memories were used most to correctly generate the answer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "selected_idx  = outputs['attentions'][28][-1] #`1971` is the 28th token\n",
    "\n",
    "topk=2 #We use topk=2 in the above example\n",
    "counter = Counter()\n",
    "stack =[]\n",
    "for layer_idx in range(32): #Look at all 32 layers, and all 32 heads\n",
    "    for head_idx in range(32):\n",
    "        stack.append(selected_idx[layer_idx].squeeze()[head_idx][-topk:].cpu().numpy())\n",
    "        counter.update(selected_idx[layer_idx].squeeze()[head_idx][-topk:].cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "counter.most_common(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks like we've got a few sequences of tokens which are retrieved frequently. Let's check out what the model found useful!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Token spans in the memory:\")\n",
    "print(tokenizer.decode(model.memory_ids.squeeze()[89:95]))\n",
    "print(tokenizer.decode(model.memory_ids.squeeze()[197:206]))\n",
    "print(tokenizer.decode(model.memory_ids.squeeze()[61:65]))\n",
    "print(tokenizer.decode(model.memory_ids.squeeze()[:6]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "innovation",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}