Personally, I’ve had to supplement my meal plan as a student with meals cooked in my dorm. With a kitchen in my room this year, I felt more ambitious in making meals, but with this came the trouble of having to deal with expiration dates, defrosting, and planning meals. I’ve trashed countless eggs after putting off making a meal with them for too long, my bacon has been waiting in the freezer for its time to shine, and I forgot about the lemons at the bottom of my fridge until they turned black. Using Foundry, I implemented a food inventory app that helps students manage their dorm food inventory, generate recipes, and minimize food waste.

This pipeline handles hundreds of food items and tons of metadata in milliseconds. In this post, I’ll walk through the project, the tools I used to accomplish it, and the challenges along the way.

The Problem

As a student, managing food inventory can be challenging:

• Buying food we already have.
• Forgetting about perishables until they expire.
• Struggling to create balanced meals with the ingredients on hand.

I had a few key solutions in mind to solve these challenges:

• Automatically track purchases via Target receipts.
• Monitor expiration dates with alerts.
• Generate AI-powered recipes using ingredients on hand.
• Reduce waste by prioritizing soon-to-expire items.

As a student at Penn State, I made the vast majority of my purchases at Target, and so I decided to implement my solution with the following:

• Target API: To fetch receipt data and items by receipt.
• Python and Pandas: To transform receipt data and create the dataset for import to Foundry.
• Palantir Foundry: For building the data pipeline, managing ontologies, and creating the app’s interface.
• Palantir AIP and GPT-4o: To generate recipes based on ingredients present and goals.

Implementation

Data Extraction and Preparation

Initially, I thought of using OCR on receipts to handle grocery data, but this was a complete nightmare. Receipts often didn’t contain sufficient useful info to determine what products were.

However, Target conveniently tracks the history of all in-store orders if you use your number for their rewards service. Logging into Target, we can see a list of purchases made:

Ideally, we would leverage a public API to get all of the items for each order, but Target has no such API. I used Chrome’s developer tools to inspect what data is being sent when navigating to the order history page, and see a request under https://api.target.com/guest_order_aggregations/v1/order_history. We can use the developer tools to inspect the request parameters as well as the returned payload to get an idea of how to use this API.

NB: I only went through this process as many times as necessary, getting info on all items ~twice. The laws for web scraping are complex, and it’s responsible to not hammer Target’s servers in this case and ethically request data. Being careful with scraping and respecting content scraped, respecting the owner of the content, and requesting data at a reasonable rate are all more important when a public API isn’t provided. Ideally, a User-Agent string would be provided with these requests, but their scope was highly limited. You can find more info on ethical scraping here.

Using an online JSON viewer, we can easily explore the returned schema from the Target API and see that receipts are returned for each purchase along with items and their thumbnails.

However, this JSON object doesn’t contain information about the quantity or price, which would be helpful in the final product. It also has the following parameters, which restrict it to only displaying the 10 most recent receipts:

page_number: 1
page_size: 10
order_purchase_type: STORE
pending_order: true
shipt_status: true

To adjust to this limitation, we use the total_pages value of the JSON object to tell how many pages to request and iterate the page_number parameter through the number of pages. I investigated the API further by opening the link to one of the orders.

Here, we see a request to https://api.target.com/guest_order_aggregations/v1/<store_receipt_id>/store_order_details, where <store_receipt_id> is the ID of the receipt from the order_history request we made earlier (in this case, 4292-3234-0163-5284). The JSON returned from this request contains all the information we need about each item, in a nested object format.

To make these API requests through a program, we need to authorize access to the API. Simply running curl on the API endpoint results in:

>curl https://api.target.com/guest_order_aggregations/v1/order_history
{"message":"Invalid Key","errors":[{"resource":"Go-Proxy","field":"","reason":"no API key found in request"}]}

Inspecting the request headers in Chrome, we can try to mimic the original request to authorize our use of the API. One of the request headers is x-api-key, which we can add to the curl request as follows:

>curl https://api.target.com/guest_order_aggregations/v1/order_history -H "x-api-key: <YOUR API KEY FROM DEV TOOLS>"
{"message":"Unauthorized","errors":[{"resource":"Go-Proxy","field":"","reason":"[GSP] No authorization header, nor auth cookies found"},{"resource":"Go-Proxy","field":"","reason":"[ID2] token value missing or Authorization invalid"}]}

Now, we see that we are still missing something, although this is likely a request cookie rather than another header. We can see the cookies sent by accessing the cookies tab in the developer tools:

My instinct was to use the accessToken cookie to authorize the request, and I sent this with the request in curl:

>curl https://api.target.com/guest_order_aggregations/v1/order_history -H "X-Api-Key: <YOUR API KEY FROM DEV TOOLS>" -H "Cookie: accessToken=<YOUR ACCESSTOKEN COOKIE>"
{"metadata":{"total_time":55,"guest_type":"P"},"orders":[],"request":{"page_number":1,"page_size":10,"shipt_status":false,"pending_order":false}}

We can see we get a valid response once this cookie is included, but we need to note that this cookie expires once the user session ends and that we would need to log in again to get a new cookie once the old one is no longer valid. However, we can ignore this for now and make requests through a simple Python program:

import requests

API_URL = 'https://api.target.com/guest_order_aggregations/v1'
ORDER_HISTORY_ENDPOINT = '/order_history'
ORDER_DETAILS_ENDPOINT = '/{store_receipt_id}/store_order_details'

PARAMS = {
    'page_size': 10,
    'order_purchase_type': 'STORE',
    'pending_order': 'true',
    'shipt_status': 'true'
}

HEADERS = {
    'X-Api-Key': '<YOUR_API_KEY>',
    'Cookie': 'accessToken=<YOUR_COOKIE>'
}

def fetch_all_receipts():
    receipts = []
    current_page = 1
    total_pages = None

    while total_pages is None or current_page <= total_pages:
        print(f'Fetching page {current_page}...')
        response = requests.get(API_URL + ORDER_HISTORY_ENDPOINT, 
                                params={**PARAMS, 'page_number': current_page}, 
                                headers=HEADERS)
        response.raise_for_status()
        data = response.json()

        orders = data.get('orders', [])
        for order in orders:
            receipts.append(order['store_receipt_id'])

        if total_pages is None:
            total_pages = data['total_pages']

        current_page += 1

    return receipts

This program will fetch the first page using the default parameters and our headers. We only get the orders list from the returned JSON, which contains the receipt ID associated with it under the store_receipt_id field. We iterate through every order to get all the receipt IDs, which we can then use to get full receipt data. If we are on the first page, we also find the total pages of orders the user has through metadata, which tells us when to stop fetching new pages.

Next, we want to fetch all items from each receipt. While data unique to each item on the receipt is stored in objects within the order_lines list, some data, such as purchase date and location purchased are only given once for the whole receipt. Because I’d like to make a dataset of items bought that contains this information by object, I make sure to include this receipt-wide data in every item’s data. The item-by-item data is also highly nested, for example, the product type information for the first item on a receipt is under response['order_lines'][0]['item']['product_classificaton']['product_type_name']. We’d like all of this information to be flattened, so we can use the Python pandas library and the pandas.json_normalize function to flatten all of this data. For consistency, we will add all of the data into a Pandas DataFrame, which makes exporting easy through pandas.to_csv.

import pandas as pd
import html

def fetch_all_items(receipts):
    items = []
    items_df = []

    for receipt in receipts:
        response = requests.get(API_URL + ORDER_DETAILS_ENDPOINT.format(store_receipt_id=receipt),
                                params={'subscription': 'false'},
                                headers=HEADERS)
        response.raise_for_status()
        data = response.json()

        items = data.get('order_lines', [])

        for item in items:
            outer_data = pd.json_normalize({
                'order_purchase_date': data['order_purchase_date'],
                'address_line1': data['address'][0]['address_line1'],
                'city': data['address'][0]['city'],
                'state': data['address'][0]['state'],
                'zip_code': data['address'][0]['zip_code'],
                'store_receipt_id': data['store_receipt_id'],
                'store_id': data['store_id'],
                'grand_total': data['grand_total']
            })
            inner_data = pd.json_normalize(item)
            inner_data['item.description'] = inner_data['item.description'].apply(html.unescape)
            items_df.append(outer_data.join(inner_data))

    return pd.concat(items_df)

In this code, we fetch each reciept, and iterate through each item’s data, flattening the “outer data” (data that applies to the whole receipt) and “inner data” (data that is by item). Then, we also unescape the item’s description tag, which often has trademark or ampersand symbols that would render as its escaped version in Palantir Foundry. Finally, we can join the inner and outer data into a single row, and append the row to a DataFrame containing each item from each receipt.

One final thing I did before importing this data to Foundry was get the image data for each item. This provided a much cleaner look in the final workshop with the icon preview for each object. Under the flattened item data, the URL for each item was accessible under image_url = 'https:' + item['item.images.base_url'] + item['item.images.primary_image']. However, this URL pointed to an AVIF image, and seemed unsupported when I tried to use it in Foundry. I also tried to use URL-based image conversion from AVIF to PNG using Cloudinary, but this also didn’t render in Foundry. As a last resort, I downloaded the image for each icon and encoded it as a base64 string, which was finally rendered in Foundry.

import io
import base64
from PIL import Image

def get_images(items):
    print('Converting images...')
    base64_images = []

    for index, item in items.iterrows():
        if type(item['item.images.base_url']) is str: # some items don't have image
            image_url = 'https:' + item['item.images.base_url'] + item['item.images.primary_image']
            
            response = requests.get(image_url)
            response.raise_for_status()

            
            image = Image.open(io.BytesIO(response.content))

            
            with io.BytesIO() as png_output:
                image.save(png_output, format='PNG')
                png_data = png_output.getvalue()

            
            base64_image = base64.b64encode(png_data).decode('utf-8')
            base64_images.append(f"data:image/png;base64,{base64_image}")
        else:
            base64_images.append('')

    items['image_base64'] = base64_images
    return items

The code above converts the AVIF image to a PNG in memory, before encoding it in base64. We can use all the functions above to generate a CSV with all order data for a user:

if __name__ == '__main__':
    receipts = fetch_all_receipts()

    items = fetch_all_items(receipts)
    items = get_images(items)

    items.to_csv('receipt_items.csv')

Data Processing in Foundry

Below is a view of the full project in Foundry. To process data in Foundry, we first upload the receipt_items dataset, process it using the Target Receipts Pipeline into the total_inventory dataset, and then display the data using the Dorm Inventory Viewer workshop and suggest recipes using the Recipe Suggestion Logic.

To process data in Foundry, we first upload the CSV file, which creates a raw dataset.

This dataset contains 43 columns, which we will drop and filter through, and all the necessary data for each item. In Foundry, we turn raw datasets into datasets usable for building programs through the Pipeline Builder.

The inputs to each pipeline are tables (datasets), either imported from Foundry or manually entered. Then, tables can be joined, unioned, or transformed. Transformation paths consist of transforms, which can be picked from a list of 100+ functions. We use these functions to filter from 43 columns to only 13.

These transforms are highly configurable and changes can easily be previewed. Each of these transforms on the path consists of a function that takes some number of columns/expressions/values and uses its output to replace or create a new column. For example, we can create a new column for the total price of some goods by multiplying the numbers in itemunit_price and quantity.

One particularly important piece of information is the purchase time and date. In the original CSV, the times appear as 2024-10-18T09:16:51-05:00, which includes the date, time, and timezone offset. From this, we can obtain timezone-specific data, but for the sake of this program, we will ignore timezones and only provide the user with local times. Foundry interprets these strings as a Struct{timestamp: Timestamp, offset: Integer}, which now has functions supporting field extraction (although these seem to be buggy right now). This is an issue when we make a Workshop interface later as Foundry offers specialized filter widgets to Timestamp and other types. To support this, I cast the structs to strings and then cast the string to a date timestamp based on a format string.

I also created a UUID column for support in the ontology later and created new columns using cast to X functions with no parameters.

Building the Ontology and AIP Logic

The previous steps in the pipeline create a dataset (as opposed to the “raw dataset” we imported) that can be used to create an Ontology using the Ontology Manager.

The Ontology uses the datasets imported into Foundry to create models of objects that represent their real-world counterparts. As an example, we use the “Item” object to represent any item bought from Target. The Ontology Manager app also allows for viewing which apps are reliant on an Ontology object, auditing changes made, creating automations, and analyzing object usage. The primary features that make the Ontology usable for operator-facing apps are the properties and actions.

The properties I decided to use for the Item object are information that may be relevant to the user managing their dorm and cooking (quantity, expiration date), as well as price and purchase date data. These properties can be backed by a dataset or stream, and I map the columns from the previously created dataset to these properties. Then, when we want to edit these objects (eg. set an expiration date, remove from inventory) we can set up actions that precisely define how users may interact with the Ontology.

As an example, I have an action that removes items from the current inventory. To maintain a history of previous items we bought for potential data analysis, I set the default empty “Removed Date” property to the current date. Then, by default, the views in the app filter out objects that have their removed date set. Actions let you create forms for these specific changes, which can be useful if you only want the user edit certain properties and validate input.

Creating the Workshop

Making the Ontology is the final stop before data can be analyzed in dozens of ways. Foundry contains several apps that can all interact with your Ontology, and Workshop allows making high-quality applications to visualize and operate on data using a point-and-click editor.

The app I created is a basic inventory manager for achieving my initial goals. Firstly, I created two different views for the items in inventory: item view and grid view. The item view uses the images generated in the preprocessing steps to generate previews, which can be very useful in managing the inventory, while the grid view allows for viewing more specific data. I then added filters that can filter on any property of the Item object and a button to remove items using the Remove Item from Inventory action. All of these components are “widgets” in Workshop, and can be dropped into the application and configured in many ways.

As an example, the “Remove Items” button was configured as shown above. Conditional visibility can be toggled using variables, which are global to the Workshop environment.

The object view, which uses checkboxes to mark selected objects, can be configured to output the selected items to a variable, which I named inventorySelectedObjects. Then, I use a “Variable Transformation” for the button visibility, which lets me define a new variable. From there, I added a transform called “Boolean comparison”, which outputs true if any conditions are matched, and set a condition that the inventorySelectedObjects object set is not empty. I can continue transforming the variable, but this is sufficient to show the button as disabled when the user hasn’t selected any objects to delete.

Finally, many things can be configured to happen when a button is clicked, but I made the button perform the action we defined earlier. Then, I can fill the form with variables from the workshop, such as inventorySelectedObjects and the current timestamp. Then, I can hide the form and immediately apply changes, since we prefilled all fields.

Separately, I added a standalone form that uses a “Create Item” action I defined earlier so that users can add their own items upon purchase. Typically, data that I showed earlier, like from the Target API, would come to Foundry as a stream, so that the datasets would be constantly updated. However, as I made this as a demo and did one-time processing (the auth cookies expire before data can be updated), I added this feature.

The final feature I added was recipe generation using AIP, Palantir’s AI platform. This widget uses the inventoryObjects variable from the complete inventory and the recipeGoals variable set by a textbox in the workshop to prompt an LLM for recipes to make.

The above is an example of the start of an AIP Logic function, which defines interactions with LLMs. These interactions can enable LLMs to act on and edit Ontology objects and are highly configurable. In this example, I decided how to format each Item into a string using a string transform block, and included the properties relevant to creating a recipe.

In the above block, I specify how an LLM is used, including the system prompt, task prompt, and strings to use. Here, I use the string variable we created and formatted earlier, alongside the user’s goals to create recipes. There are over 200 different blocks that can be used in this AIP Logic, which allows for highly custom LLM and tool usage. However, I simply used the LLM output from this in my Workshop.

Dropping in an AIP Logic widget, I can let the user generate recipes using the goal provided and items in inventory. I initially struggled as the block itself has many configurations, and used the AIP Logic as input to another LLM call, which resulted in the output being an LLM response to the first LLM response. Upon fixing this bug, the widget instantly lets me generate recipes and provides suggestions for food to buy to support my metabolic health, or whatever my goals are.

Future Improvements

I strongly believe Foundry is a revolutionary product, and that I would never be able to create an application as high-quality in a short amount of time. Although Foundry has a steep learning curve, it is a powerful tool for automating decision-making. There are dozens of apps I haven’t explored in this small example, including Quiver, which can be used to easily analyze data:

Although I made this app in the span of 3-4 days through on-and-off work, I would add several features to make it work best for me. Firstly, I had hoped to use vision models to be able to generate Item objects from grocery receipts but didn’t do it in the time span of the project. Barcodes were another potential option, although this also was difficult to do with Foundry in its current state.

Additionally, I would use the Ontology features embedded in AIP Logics to generate recipe objects that can be saved and used (to automatically remove items from inventory). I could integrate Fitbit data in the future to determine how my diet is affecting my health, and use that data alongside spending to build a budget-friendly grocery list using AIP. Other data, such as item expiry data, could be automatically implemented based on online databases, and alerts could be provided alongside recipes to avoid friction in manually entering these dates.

Conclusion

Foundry is incredibly powerful, and this demo only scratches the surface of what is possible. Palantir’s demo of Foundry from 2022 showcases many more features that I didn’t have the chance to use.

Could this app cut food waste among students in half? I’ve noticed that it was much easier for me to figure out what to do with my food once I built this. But this would still be one of the less impactful things made with Foundry, as it is used by dozens of clients to automate decision making across the board.

I’m excited to ship new features for Foundry during my internship at Palantir and hope that others find use out for powerful software.

Technologies and Integrations

SweetFriend utilizes an array of technologies and external APIs to deliver a seamless user experience:

• We used React for building the frontend user interfaces, enabling dynamic and interactive graph visualizations for glucose levels.
• Flask was chosen for the backend, serving the API and handling the business logic.
• For data storage, we opted for MongoDB due to its flexibility in schema design and efficient handling of user logs and health data.
• One of the standout features of SweetFriend is its integration with the Dexcom API, allowing users to import real-time glucose data from their Dexcom devices. This integration is vital for providing continuous glucose monitoring and timely insights.
• To enhance the AI capabilities, we leveraged Tune Studio for vision model inference using GPT-4o, estimating carbohydrate content from food images.
• Cerebras was utilized for super-fast AI chatbot inference, enabling natural and responsive conversations with users.
• For notifications and alerts, we integrated Twilio, ensuring timely communication about critical health metrics like glucose levels.

Technical Challenges and Solutions

Environment Configuration

We needed to manage configurations for different environments, including local development and Vercel deployment. The Flask application checks if it’s running on Vercel using an environment variable and adjusts its configuration accordingly:

if os.environ.get('VERCEL', None) != "True":
    app.config.from_file('config.json', load=json.load)
else:
    # Read configurations from environment variables
    app.config['TWILIO_AUTH_TOKEN'] = os.environ['TWILIO_AUTH_TOKEN']
    # Other environment-specific configs

User Authentication and Data Security

We implemented user authentication using hashed passwords and sessions. Flask’s check_password_hash and generate_password_hash functions ensure that user passwords are securely stored and verified.

from werkzeug.security import check_password_hash, generate_password_hash

@app.route('/api/user_login', methods=['POST'])
def user_login():
    username = request.form['username']
    password = request.form['password']
    user = db.users.find_one({'username': username})
    if user and check_password_hash(user['password'], password):
        session['username'] = username
        session['user_id'] = str(user['_id'])
        session['logged_in'] = True
        return redirect('/app/dashboard')
    else:
        return redirect('/')

Efficient and Secure API Endpoints

To ensure efficient and secure data retrieval, we implemented caching with Flask-Caching and used environment variables for API credentials. The Cache class also ensures that LLM outputs are cached to avoid requesting new guidance every time the user visits their dashboard, saving on API costs and serving time.

from flask_caching import Cache
cache = Cache(config={'CACHE_TYPE': 'simple'})
cache.init_app(app)

@app.route('/api/get_advice')
@cache.cached(timeout=60)
def get_advice():
    recent_glucose, recent_events, recent_logs = get_recent_data()

    recent_glucose = [{'time': x['system_time'], 'glucose_value': x['glucose_value']} for x in recent_glucose]
    recent_events = [{'time': x['system_time'], 'event_type': x['event_type'], 'value': x['value']} for x in recent_events]

    context = f"Recent glucose readings (mg/dL): {recent_glucose}\nRecent events: {recent_events}\nRecent logs: {recent_logs}\n\n"

    # ... get response based on context from Cerebras

Image Analysis for Meal Logging

We developed an endpoint for analyzing meal images using AI models. This is a significant improvement over existing glucose management apps, which require manually entering the carbohydrate values of meals the user has eaten. This causes friction in the management process which can lead to the user entirely missing or delaying meal entry. In addition, users aren’t always accurately entering the carbohydrate value of their meals. To use GPT-4o to analyze images, the uploaded image is converted to a base64 string and sent to the Tune Studio API for inference:

@app.route('/api/analyze_image', methods=['POST'])
def analyze_image():
    if 'file' not in request.files:
        return jsonify({"error": "No file uploaded"}), 400
    image = request.files['file']
    base64_image = base64.b64encode(image.read()).decode('utf-8')
    # Send to Tune Studio for analysis and return results

We found that performance improves when giving the model a chance to think about what’s in the image, then to return a JSON object of the carb values. We first analyze and break down the image using this system and user prompt:

SYSTEM:
You do not use markdown headers or bolding, just lists. You are helping diabetes patients estimate the number of carbs in their meal so they can plan and monitor their glucose levels accordingly.

USER:
First, analyze the image and describe what food items are present. Then, break down the ingredients and estimate the carbs of each ingredient in grams, then calculate the total carbs and give the name of the meal.

Then, the output of this is appended to the current conversation log. Then, we prompt the model again using this message:

Using your estimates, fill the structured output JSON with the values. Only output valid JSON according to the schema. Only output for the meal total, one entry. Do not use code blocks or anything to surround json. Meal carbs must be in grams. Write in following format:
{
    "meal_name": {"type": "string"},
    "total_carbs": {"type": "number"}
}

In Tune Studio’s API, we can prompt the model to answer in this format using the "guided_json" parameter in the API request:

"guided_json": {
    "type": "object",
    "properties": {
        "meal_name": {"type": "string"},
        "total_carbs": {"type": "number"}
    }
},

Then, we return the JSON to populate the fields in the meal entry interface.

Real-Time Notifications

Using Twilio, we set up real-time notifications to alert users about critical glucose levels:

@app.route('/api/twilio_send')
def twilio_send():
    message = twilio_client.messages.create(
        from_='number',
        body='SweetFriend: Your glucose level is low and falling. Have a snack with around 15-20g of carbs.',
        to='number'
    )
    return jsonify({'status': 'message sent', 'message_sid': message.sid})

When fully developed, our app would automatically notify the user with actionable steps when their glucose levels have fallen to dangerously low or high levels. In the future, this could also be integrated with remote patient monitoring to alert healthcare providers or caretakers when action must be taken.

Frontend Development

We built an intuitive user interface using React, incorporating various components for registration, login, and dashboard functionalities. Below are some of the highlights:

Dynamic Dashboard

The dashboard displays glucose readings and AI-generated suggestions, updating in real-time. We used useEffect to fetch data and update the state:

import { useState, useEffect } from "react";

function Dashboard() {
    const [glucoseReadings, setGlucoseReadings] = useState([]);
	const [aiSuggestion, setAiSuggestion] = useState("Random AI suggestions will be displayed here...");

    useEffect(() => {
        const fetchGlucoseData = async () => {
            try {
                const response = await fetch("/api/get_glucose");
                const data = await response.json();
                setGlucoseReadings(data);
            } catch (error) {
                console.error("Error fetching glucose data:", error);
            }
        };

        const fetchAdvice = async () => {
            try {
                const response = await fetch("/api/get_advice");
                const data = await response.json();
                setAiSuggestion(data.response);
            } catch (error) {
                console.error("Error fetching AI advice:", error);
            }
        };

        fetchGlucoseData();
        fetchAdvice();
    }, []);

    return (
        // simplified dashboard
        <div>
            <h1>Dashboard</h1>
            <div>
                {glucoseReadings.map((reading, index) => (
                    <p key={index}>Time: {reading.systemTime}, Value: {reading.value}</p>
                ))}
            </div>
            <p>{aiSuggestion}</p>
        </div>
    );
};

Using Chart.js, we visualize the glucose data alongside the logged user events so the user can see associations between exercise, meals, and glucose level. This is another improvement upon current glucose monitoring apps, which only show the uncontextualized glucose level.

We also added a feature that allows the user to click on the graph to add logs directly from the dashboard for the time selected, further simplifying the data entry process.

The past 24 hours of user data is also filtered and passed to the Cerebras API, using Llama-3.1-70b. This provides instant results that are more convient for the user. The LLM is prompted to avoid giving medical advice, and to instead give general actionable steps or insights based on recent data.

Finally, the most recent logs are also available through the dashboard, allowing users to get quick access to change their logs or star them. Starring logs allows users to mark logs that may be anomalous or that they wish to report to their healthcare provider, as they can filter their logs to instantly get a list of starred logs.

Image Upload for Meal Logging

Logging carbohydrate intake and exercise is critical for managing glucose levels. We let users do that easily through our log page:

The page has the ability to search logs, filter, and sort them. We implemented this on the front end after querying for all the logs once:

function Log() {
    // Handle sorting items
	const sortItems = () => {
		let sortedItems = [...items];
		if (sortOrder === "nameAsc") {
			sortedItems = sortedItems.sort((a, b) => a.name.localeCompare(b.name));
		} else if (sortOrder === "nameDesc") {
			sortedItems = sortedItems.sort((a, b) => b.name.localeCompare(a.name));
		} else if (sortOrder === "dateAsc") {
			sortedItems = sortedItems.sort((a, b) => a.timestamp - b.timestamp);
		} else if (sortOrder === "dateDesc") {
			sortedItems = sortedItems.sort((a, b) => b.timestamp - a.timestamp);
		}
		setItems(sortedItems);
		setShowSort(false); // Close sort dialog after sorting
	};

	// Handle filtering items by type or starred
	const applyFilters = () => {
		let filteredItems = allItems;

		if (!filters.food) {
			filteredItems = filteredItems.filter(
				(item) => item.type.toLowerCase() !== "food"
			);
		}
		if (!filters.exercise) {
			filteredItems = filteredItems.filter(
				(item) => item.type.toLowerCase() !== "exercise"
			);
		}
		if (filters.starred) {
			filteredItems = filteredItems.filter((item) => item.starred);
		}

		setItems(filteredItems);
		setShowFilters(false); // Close filter dialog after applying filters
	};

	// Handle checkbox filter changes
	const handleFilterChange = (e) => {
		const { name, checked } = e.target;
		setFilters((prevFilters) => ({ ...prevFilters, [name]: checked }));
	};

	// Handle radio button sort changes
	const handleSortChange = (e) => {
		setSortOrder(e.target.value);
	};
}

The image upload feature lets users log meals by using LLMs to analyze photos of their food. Here’s a snippet of the image upload component that posts to the /api/analyze_image endpoint shown earlier:

import { useState, useEffect } from "react";

const handleImageUpload = async (e) => {
    e.preventDefault(); // Prevent default form submission
    setIsLoading(true); // Start loading state

    const formData = new FormData(e.target);

    try {
        const response = await fetch("/api/analyze_image", {
            method: "POST",
            body: formData,
        });
        const data = await response.json();

        // Assuming API returns { meal_name, total_carbs, reason }
        setName(data.meal_name || ""); // Set the name from AI response
        setTotalCarbs(data.total_carbs || ""); // Set the total carbs from AI response
        setReason(data.reason || ""); // Set the reasoning from AI response
        setIsLoading(false); // Stop loading state
    } catch (error) {
        console.error("Error uploading image:", error);
        setIsLoading(false); // Stop loading on error
    }
};

Once the image is uploaded, the modal is given the LLM-generated description of the food and is populated with the values returned in the JSON object. This gives the user reasoning on how carbs are calculated for each ingredient, which makes it easy to adjust for major errors. The name of the meal is also populated and editable.

Deployment and Configuration

Vercel Deployment

For seamless deployment, we configured Vercel to handle both the frontend and backend:

{
    "rewrites": [
        { "source": "/api/(.*)", "destination": "/api/app" },
        { "source": "/login", "destination": "/api/app" },
        { "source": "/callback", "destination": "/api/app" },
        { "source": "/(.*)", "destination": "/index.html" }
    ]
}

The generated React website (at /index.html) is the destination for all paths, except for API-related paths, whose destination is /api/app.

Local Development

For local development, we used Vite’s proxy feature to route API calls to the Flask server running on localhost:5000:

// vite.config.js
export default {
    server: {
        proxy: {
            '/api': {
                target: 'http://localhost:5000',
                changeOrigin: true,
            }
        }
    }
}

Conclusion

SweetFriend is a significant step forward in diabetes management, combining real-time data from Dexcom devices with user input. The user is able to authenticate usage of the Dexcom API with Sweetfriend, allowing for real-time glucose data to appear in the dashboard. Logging fitness and carb intake is made easy using AI-powered image recognition. The user sees all data at their dashboard, alongside general advice on how they are managing their glucose throughout the day. An additional chatbot page allows the user to get more specific advice and ask questions, such as how to prepare future meals or plan their fitness.

Future steps for SweetFriend would include fitness tracker integration to make activity logging even easier, remote patient monitoring to enhance patient outcomes, and improved user experience through notifications, mobile app, etc.

Below is a demo of the website at the time of the hackathon, feel free to also try it here using username and password admin.

Principal Component Analysis (PCA) is an important part of machine learning and data science, as it allows you to compress data into fewer dimensions. Later in this post, we see images that were originally in 3072 dimensions being represented in just two. Despite this, major patterns and trends are “summarized” into features such that, even with only two dimensions, we can predict with higher accuracy what is in an image.

Objective

The objective of this notebook is to implement Principal Component Analysis (PCA) and Dual PCA on the CIFAR-10 dataset and compare their computational efficiency by measuring the time taken to compute the principal components.

Dataset Preparation

To implement PCA and Dual PCA, we will be using numpy. For visualizations of the results of PCA and Dual PCA, we use matplotlib. The CIFAR-10 dataset will be loaded in via the torchvision library.

The CIFAR-10 dataset contains 32x32 color images from 10 classes. There are 50,000 training images and 10,000 test images.

import numpy as np
import torch
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

# transform to torch.Tensor to later transform to numpy arrays
tensor_transform = transforms.Compose([
    transforms.ToTensor(),
])

# load CIFAR-10
training_data = datasets.CIFAR10(
    root="CIFAR-10",
    train=True,
    download=True,
    transform=tensor_transform
)

Files already downloaded and verified

# transform images from tensors to numpy
images = np.array([np.array(image[0]) for image in training_data])

images.shape

(50000, 3, 32, 32)

Now that we loaded the CIFAR-10 dataset, we will center the data by subtracting the mean from each feature. However, we first need to shape the data from a tensor into a matrix. The tensor currently has dimension $N\times C\times D \times D$, where $N$ is the number of data points, $C$ is the number of color channels, and $D$ is the dimension of the image. To turn this into a matrix, we combine the final three dimensions into a single dimension: \(N\times C \times D\times D \implies (CD^2)\times N\)

This gives us matrix $X \in \mathbb{R}^{d\times n}$. We let numpy handle reshaping by using numpy.reshape.

N, C, D, _ = images.shape

images = np.reshape(images, (N, C * D * D)) # reshape first, then transpose to maintain data order
images = images.T
images.shape

(3072, 50000)

PCA Implementation

Standard PCA requires us to compute the covariance matrix of our dataset $X\in\mathbb{R}^{d\times n}$, where $d$ is the number of features and $n$ is the number of data points. We compute the covariance matrix as follows: \(C=\frac{1}{n-1}XX^T\) First, we need to centralize the data by subtracting the mean of each feature. Then, we use the formula above to find the covariance matrix.

# get mean of each feature and subtract
X = images - np.mean(images, axis=1)[:,None] # None used to create axis to broadcast (taken from numpy.cov)

C = X @ X.T / (N - 1)

Now, we perform the eigenvalue decomposition on the covariance matrix to obtain the principal components. For this we use numpy.linalg.eig.

# find eigendecomposition
eigvals, eigvecs = np.linalg.eig(C)

# we sort the eigenvectors by the eigenvalues to order the principal components
idx = np.argsort(-eigvals, axis=0)
PCvals = eigvals[idx]
PCvecs = eigvecs[:, idx]

Let’s implement this as a function from start to finish to profile the performance of later.

def pca(images, num_images):
    # construct matrix X and centralize
    X = images[:,:num_images]
    X = X - np.mean(X, axis=1)[:,None]

    # compute the covariance matrix
    C = X @ X.T / (num_images - 1)
    
    # perform eigenvalue decomposition and sort by principal components
    eigvals, eigvecs = np.linalg.eig(C)
    idx = np.argsort(-eigvals, axis=0)
    PCvals = eigvals[idx]
    PCvecs = eigvecs[:, idx]
    
    return PCvals, PCvecs

Dual PCA Implementation

To compute the dual PCA, we use the same matrix $X$, but subtract the mean from each data point instead. Then, we compute the Gram matrix $G$: \(G=\frac{1}{n-1}X^TX\) Then, we perform eigenvalue decomposition on the $n\times n$ Gram matrix to find the principal components.

X = images - np.mean(images, axis=0)
X = X[:,:5000]

G = X.T @ X / (4999)

Then, we perform eigenvalue decomposition to obtain the principal components. Note that the encoding and reconstruction of training data is different than with standard PCA.

Encoding

We encode the training data by first calculating the Gram matrix and finding the eigenvalue decomposition. Then, $V\in\mathbb{R}^{n\times p}$ is the matrix of the top $p$ eigenvectors of $G$ and $\Sigma\in\mathbb{R}^{p\times p}$ is a diagonal matrix of the square roots of the top $p$ eigenvalues. The encoding of the data $Y\in\mathbb{R}^{p\times n}$ is:

For an individual test sample, we can get a $p$-dimensional encoding $\mathbf{y}$ of $\mathbf{x}$:

Reconstruction

The reconstructed training data $\hat X\in\mathbb{R}^{d\times n}$ can be obtained as shown:

To reconstruct a test sample $\hat{\mathbf{x}}$, we use the formula below:

eigvals, eigvecs = np.linalg.eig(G)

idx = np.argsort(-eigvals, axis=0)
DPCvals = eigvals[idx]
DPCvecs = eigvecs[:, idx]

From the cell execution times above, we can see that the dual PCA algorithm took much longer to run in comparison to the standard PCA implementation. This is because we are doing a matrix multiplication along the larger dimension ($n > d$). In addition, the large result (50000 x 50000 for the full dataset) cannot fit within most memory sizes, also making computation difficult. Therefore, we reduce the dataset to the first 5000 images in this example.

We will implement a function to find the relevant Dual PCA vectors for further profiling.

def dual_pca(images, num_images):
    # construct matrix X and centralize
    X = images[:,:num_images]
    X = X - np.mean(X, axis=1)[:,None]

    # compute the Gram matrix
    G = X.T @ X / (num_images - 1)
    
    # perform eigenvalue decomposition and sort by principal components
    eigvals, eigvecs = np.linalg.eig(G)
    idx = np.argsort(-eigvals, axis=0)
    PCvals = eigvals[idx]
    PCvecs = eigvecs[:, idx]
    
    return DPCvals, DPCvecs

Methods Comparison

We will measure the time taken to get principal components by PCA and Dual PCA on the CIFAR-10 dataset, recording for different data sizes/dimensions.

Below we have a series of dataset sizes to test, and we take the first $n$ images in each case for PCA.

import time

# different choices of n to test
pca_image_sizes = [100, 500, 1000, 2500, 5000, 10000, 20000, 30000, 40000, 50000]
pca_time_taken = {}

# time all configurations of n
for size in pca_image_sizes:
    start_time = time.time()
    pca(images, size)
    end_time = time.time()
    pca_time_taken[size] = end_time - start_time

for size, duration in pca_time_taken.items():
    print(f"Time taken for PCA with {size} images: {duration:.4f} seconds") 

Time taken for PCA with 100 images: 30.5286 seconds
Time taken for PCA with 500 images: 27.9196 seconds
Time taken for PCA with 1000 images: 24.4410 seconds
Time taken for PCA with 2500 images: 22.2044 seconds
Time taken for PCA with 5000 images: 24.0145 seconds
Time taken for PCA with 10000 images: 24.5102 seconds
Time taken for PCA with 20000 images: 25.3693 seconds
Time taken for PCA with 30000 images: 26.5230 seconds
Time taken for PCA with 40000 images: 27.4280 seconds
Time taken for PCA with 50000 images: 27.3944 seconds

# only time up to n = 10000 images to avoid costly computation and OOM error
dpca_image_sizes = [100, 500, 1000, 2500, 5000, 10000]
dpca_time_taken = {}

for size in dpca_image_sizes:
    start_time = time.time()
    dual_pca(images, size)
    end_time = time.time()
    dpca_time_taken[size] = end_time - start_time

for size, duration in dpca_time_taken.items():
    print(f"Time taken for Dual PCA with {size} images: {duration:.4f} seconds") 

Time taken for Dual PCA with 100 images: 0.0215 seconds
Time taken for Dual PCA with 500 images: 0.5890 seconds
Time taken for Dual PCA with 1000 images: 1.8850 seconds
Time taken for Dual PCA with 2500 images: 10.8569 seconds
Time taken for Dual PCA with 5000 images: 57.1203 seconds
Time taken for Dual PCA with 10000 images: 343.6070 seconds

We can see that the time for the PCA algorithm remains constant, while the Dual PCA time increases polynomially. Assuming $O(n^3)$ time complexity for np.linalg.eig, modeling the execution time in Desmos ($y_1\sim ax_1^3$) for the full dataset produces an estimate of over 14 hours runtime. This would also require a lot of memory, which makes the calculation infeasible.

plt.figure(figsize=(12, 8))
ax = plt.subplot(111)
ax.bar(np.arange(len(pca_time_taken))-0.15, pca_time_taken.values(), width=0.3, color='b', align='center', label="PCA")
ax.bar(np.arange(len(dpca_time_taken))+0.15, dpca_time_taken.values(), width=0.3, color='r', align='center', label="Dual PCA")
ax.set_xticks(np.arange(len(pca_time_taken)))
ax.set_xticklabels(pca_time_taken.keys())
ax.set_ylabel("Time (s)")
ax.set_xlabel("Number of images")
ax.legend()

plt.show()

We have no values for the execution time of the dual PCA algorithm at $n>10000$ because the time to execute and memory requirement increases significantly. The performance of the PCA algorithms is primarily dependent on the eigenvalue decomposition, rather than the matrix multiplication to find the covariance or Gram matrix. This means that the PCA performance is roughly constant because we are finding the decomposition for a $d \times d$ matrix, and $d$ is constant in these tests. The dual PCA execution time varies with $n$, and is significantly faster when $n < d$, as expected.

Therefore, we can conclude that dual PCA is advantageous for high-dimensional, small datasets where $n < d$. Otherwise, standard PCA will execute at about the same speed for datasets as they grow in size.

Visualization

We will begin by visualizing the principal components and their values. Below we plot the sorted principal components and their associated eigenvalues.

plt.figure(figsize=(16, 6))

plt.subplot(121)
plt.yscale('log')
plt.plot(range(1,len(PCvals)+1),PCvals,'.')
plt.title("PCA")
plt.ylabel('Eigenvalues of covariance matrix',fontsize=15)
plt.tick_params(labelsize=13)

plt.subplot(122)
plt.yscale('log')
plt.plot(range(1,len(DPCvals)+1),DPCvals,'.')
plt.title("Dual PCA")
plt.ylabel('Eigenvalues of Gram matrix (n=5000)',fontsize=15)
plt.tick_params(labelsize=13)

plt.show()

We can see that the first few principal components are the most “important”, and that there is then a significant drop in the eigenvalues before it only decreases exponentially. On the dual PCA side, we can see that there is a sharp drop in the eigenvalues at $n=3072$, which corresponds to the dimension of the images in the dataset.

import cv2

pc_inds = [1,2,5,10,15,20,50,100,150,200,500,1000]

plt.figure(figsize=(12, 10))

for i in range(0, len(pc_inds)):
    plt.subplot(3, 4, i + 1)

    ind = pc_inds[i]
    
    # seperate principal component into RGB color channels and normalize
    R, Gr, B = PCvecs[:, ind - 1].reshape(3, 32, 32)
    R = (R - R.min()) / (R.max() - R.min()) * 255
    Gr = (Gr - Gr.min()) / (Gr.max() - Gr.min()) * 255
    B = (B - B.min()) / (B.max() - B.min()) * 255

    # merge RGB channels to create image
    im = cv2.merge([R, Gr, B]).astype(np.uint8)
    plt.imshow(im)

    plt.axis('off')
    plt.title('PC = {}'.format(ind))

plt.tight_layout()
plt.show()

We can visualize the principal components by reshaping the eigenvectors back to a 3-channel 32x32 image. By combining these images as specified by an encoded image, we can reconstruct an approximation of the original data.

The principal components are harder to visualize for the dual PCA method, as the eigenvectors are of dimension $\mathbb{R}^n$, not $\mathbb{R}^d$.

Below we show how to reconstruct an image using the PCA method and how the image quality varies with the number of principal components $p$ used.

plt.figure(figsize=(12, 10))
plt.suptitle("PCA Reconstruction")

for i in range(0, len(pc_inds)):
    plt.subplot(3, 4, i + 1)

    ind = pc_inds[i]
    
    # reconstruct image xhat = UU^Tx
    xhat = PCvecs[:,:ind] @ PCvecs[:,:ind].T @ X[:,1]
    
    R, Gr, B = xhat.reshape(3, 32, 32)
    R = (R - R.min()) / (R.max() - R.min()) * 255
    Gr = (Gr - Gr.min()) / (Gr.max() - Gr.min()) * 255
    B = (B - B.min()) / (B.max() - B.min()) * 255

    im = cv2.merge([R, Gr, B]).astype(np.uint8)
    plt.imshow(im)
    
    plt.axis('off')
    plt.title('PC = {}'.format(ind))

plt.tight_layout()
plt.show()

Above we reconstruct the images using the PCs of the standard PCA algorithm. As the number of components used increases, the image is a closer approximation of the original image. We can also see that the first principal components have the largest impact on the final image, while at $p=1000$, the image is already close to the full image.

plt.figure(figsize=(12, 10))
plt.suptitle("Dual PCA Reconstruction")

for i in range(0, len(pc_inds)):
    plt.subplot(3, 4, i + 1)
    
    ind = pc_inds[i]
    
    # reconstruct xhat = XVS^-2V^TX^Tx
    xhat = X @ DPCvecs[:,:ind] @ np.linalg.inv(np.diag(DPCvals[:ind])) @ DPCvecs[:,:ind].T @ X.T @ X[:,1]
    
    R, Gr, B = xhat.reshape(3, 32, 32)
    R = (R - R.min()) / (R.max() - R.min()) * 255
    Gr = (Gr - Gr.min()) / (Gr.max() - Gr.min()) * 255
    B = (B - B.min()) / (B.max() - B.min()) * 255
    
    im = cv2.merge([R, Gr, B]).astype(np.uint8)
    plt.imshow(im)
    
    plt.axis('off')
    plt.title('PC = {}'.format(ind))

plt.tight_layout()
plt.show()

Above we reconstruct the image using dual PCA. We achieve similar results, with no significant difference between this and the standard PCA reconstruction.

With the principal components, we can remap the 3072-dimensional image space into two dimensions (using the first two principal components). We encode each image using the principal component vector and then plot each image in 2D.

proj = np.dot(images.T, PCvecs[:,:2])
plt.scatter(proj[:,0], proj[:,1], c=training_data.targets, s=2)
plt.title('PCA Projection of CIFAR-10 Images')
plt.xlabel('Principal Component 1')
plt.ylabel('Principal Component 2')
cbar = plt.colorbar()
cbar.set_ticks(np.arange(10))
cbar.set_ticklabels(training_data.classes)
plt.show()

In the above plot, we also colored each point using the target value of the image, which corresponds to its class (truck, plane, etc.). The grouping of the various classes shows that PCA, in capturing variance, also preserves some of the features that are the same among different images of the same class. This is true despite the fact that, as we saw earlier, the first principal components only capture very general lighting conditions of the image.

from mpl_toolkits.mplot3d import Axes3D

proj_3d = np.dot(images.T, PCvecs[:,:3])
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
scatter = ax.scatter(proj_3d[:,0], proj_3d[:,1], proj_3d[:,2], c=training_data.targets, s=1)
ax.set_title('3D PCA Projection of CIFAR-10 Images')
ax.set_xlabel('Principal Component 1')
ax.set_ylabel('Principal Component 2')
ax.set_zlabel('Principal Component 3')

# rotate the view
ax.view_init(45, 45)

plt.show()

In the above graph, we plot the third principal component as another dimension. The strength of the grouping is marginally better than the graph with only two PCs, and we can expect subsequent dimensions to have a diminishing effectiveness in separating images of the same class into clusters.

Next, we encode the full image dataset using the first two principal components from the dual PCA algorithm.

# encode with Xhat = S^-1V^TX^TX
proj = np.sqrt(np.linalg.inv(np.diag(DPCvals))) @ DPCvecs.T @ X.T @ images
plt.scatter(proj[0,:], proj[1,:], c=training_data.targets, s=5)
plt.title('PCA Projection of CIFAR-10 Images')
plt.xlabel('Principal Component 1')
plt.ylabel('Principal Component 2')
plt.show()

Although this is different than the result from the PCA plot, this is to be expected as dual PCA is only performed on the first 5000 images. We can also notice that despite being outside of the training data, the dual PCA matrix also effectively encodes images beyond the first 5000 in a way that brings similar classes to close proximity.

Conclusion

PCA and dual PCA are both important algorithms that can be used to transform higher-dimensional data to a lower-dimensional space. PCA involves computing a $\mathbb{R}^{d\times d}$ matrix, which is expensive for high-dimensional data, such as images. By using dual PCA when $d > n$, we can save on computation by computing a $\mathbb{R}^{n\times n}$ matrix instead. As seen with CIFAR-10 above, dual PCA can be highly effective as $d$ grows relative to $n$, as a 4k image would already have over 24 million dimensions.

Security Enhancements

Polyfill Fix

polyfill.io is used for loading polyfills (code that provides modern functionality on older browsers). Mathjax, the library MarkWeb uses for rendering mathematical equations, relies on polyfills for compatibility with older browsers.

In light of the supply chain attack on polyfill.io, the link to polyfill.io has been replaced with a link to Cloudflare’s mirror. This change ensures that the polyfills are served securely and reliably.

old polyfill.io link:

<script src="https://polyfill.io/v3/polyfill.min.js?features=es6"></script>

new Cloudflare mirror link:

<script src="https://cdnjs.cloudflare.com/polyfill/v3/polyfill.min.js?version=4.8.0&features=es6"></script>

Proxy Security

Another security feature coming eventually is rate limiting. This is typically done by using the remote_address (client’s IP address). However, as mentioned in the earlier blog, Cloudflare is used as a proxy for our app, which means that the client’s IP will appear to be the Cloudflare proxy server. To determine the proper client IP address, the X-Forwarded-For header is used. For each proxy, an IP address is concatenated to the X-Forwarded-For header.

X-Forwarded-For: <client>, <proxy1>, <proxy2>

Because we trust Cloudflare as a proxy, <proxy1> should equal a Cloudflare proxy IP address, while <client> should be the remote_address we use to rate limit. However, this header can be easily spoofed by attackers, leading to a false IP address being used for access control. For example, we can use curl to add a proxy IP to be forwarded to the Cloudflare proxy. In the example below, we use a special route that returns headers and request.remote_addr, among other info.

curl -H "X-Forwarded-For: 127.0.0.1, 1.1.1.1" http://markweb.test:5000/debug/test_proxy
{
  "CF-Connecting-IP": null,
  "Client IP": "127.0.0.1",
  "Host": "markweb.test:5000",
  "Remote Addr": "127.0.0.1",
  "Scheme": "http",
  "X-Forwarded-For": "127.0.0.1, 1.1.1.1, 2.2.2.2",
  "X-Forwarded-Host": null,
  "X-Forwarded-Port": null,
  "X-Forwarded-Prefix": null,
  "X-Forwarded-Proto": null
}

In the above scenario, all requests are passed through a proxy with IP address 2.2.2.2. The true client IP is 1.1.1.1, but we set the X-Forwarded-For header to 127.0.0.1, 1.1.1.1. By default, Flask will blindly trust the first X-Forwarded-For header value as the true client IP address, setting request.remote_addr to this value. We successfully spoofed the request, which lets us bypass any IP-based rate limiting or access IP restricted routes and features.

To avoid this, we let the Flask app know how many proxies to trust the header for. This is done with Werkzeug’s ProxyFix middleware. For a single proxy like Cloudflare, we setup the middleware as follows:

from werkzeug.middleware.proxy_fix import ProxyFix

# ...

app.wsgi_app = ProxyFix(app.wsgi_app, x_for=2, x_proto=1, x_host=1, x_port=1, x_prefix=1)

This trusts the second to last value of the X-Forwarded-For header, or the one right before the last proxy.

curl -H "X-Forwarded-For: 127.0.0.1, 1.1.1.1" http://markweb.test:5000/debug/test_proxy
{
  "CF-Connecting-IP": null,
  "Client IP": "127.0.0.1",
  "Host": "markweb.test:5000",
  "Remote Addr": "1.1.1.1",
  "Scheme": "http",
  "X-Forwarded-For": "127.0.0.1, 1.1.1.1, 2.2.2.2",
  "X-Forwarded-Host": null,
  "X-Forwarded-Port": null,
  "X-Forwarded-Prefix": null,
  "X-Forwarded-Proto": null
}

Now, the client IP field (determined from X-Forwarded-For) is the untrusted first value. The remote address field, which is used in the application for IP-related features, is properly set to 1.1.1.1, based on the number of values to trust. The rest of the X-Forwarded header values should only be trusted for 1 value (set by Cloudflare proxy), and so we set those parameters to 1. More info on the ProxyFix middleware is in Werkzeug’s documentation.

You may also notice the CF-Connecting-IP header is printed in our debug output. This is a non-standard header that Cloudflare uses to immediately communicate the client IP. Because we are only using one proxy, we can write code more specific to Cloudflare’s proxy service as follows:

class CloudflareProxyFix:
    def __init__(self, app):
        self.app = app

    def __call__(self, environ, start_response):
        environ_get = environ.get
        environ["cloudflare_proxy_fix.orig"] = { 
            "REMOTE_ADDR": environ_get("REMOTE_ADDR") 
        }
        cf_connecting_ip = environ_get('HTTP_CF_CONNECTING_IP')
        if cf_connecting_ip:
            environ["REMOTE_ADDR"] = cf_connecting_ip

        return self.app(environ, start_response)

Following the conventions of the original ProxyFix, we set environ_get to environ.get for efficiency, and then update the environment with the original value of REMOTE_ADDR, for debugging and logging purposes. Then, if the CF-Connecting-IP header is present, we set the remote address to it. Because CF-Connecting-IP is always set by Cloudflare and any client value is ignored, we can generally avoid implementing extra checks for this header.

Performance Optimizations

Efficient Static File Serving

One of the most useful things to do when first optimizing a web application is to profile it. To start profiling, I added a profiling config flag to the app. When this flag is set, the app will use the Werkzeug profiler middleware to profile the app.

if app.config['PROFILER']:
    app.wsgi_app = ProfilerMiddleware(app.wsgi_app, restrictions=[0.1])

With restrictions set to 0.1, this middleware prints the top 10% of the slowest function calls, sorted by total time taken.

--------------------------------------------------------------------------------
PATH: '/'
         2471 function calls (2434 primitive calls) in 0.874 seconds

   Ordered by: internal time, call count
   List reduced from 737 to 74 due to restriction <0.1>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        2    0.662    0.331    0.662    0.331 {method 'execute' of 'psycopg2.extensions.cursor' objects}
        2    0.201    0.100    0.201    0.100 {method 'rollback' of 'psycopg2.extensions.connection' objects}
        4    0.002    0.000    0.002    0.000 {method 'write' of '_io.TextIOWrapper' objects}
        1    0.000    0.000    0.000    0.000 {built-in method nt.stat}
      195    0.000    0.000    0.000    0.000 {built-in method builtins.isinstance}
      3/1    0.000    0.000    0.000    0.000 c:\users\denni\markweb\.venv\lib\site-packages\sqlalchemy\sql\cache_key.py:221(_gen_cache_key)
      ...

In addition to showing the time taken for Flask to serve the path requested, it also shows the profile of all other functions called, including those serving favicon.ico and other static files. One thing I noticed when looking at the static file profiles was that database queries were being made for each static file.

This was because the before_app_request decorator was being used to load the logged-in user. This decorator is called for every request, including those for static files. To avoid this, I added a check to skip the user loading for static files and other non-essential paths.

@bp.before_app_request
def load_logged_in_user():
    if request.path.startswith('/static') or request.path in [
        '/favicon.ico', '/robots.txt', 
        '/sitemap.xml', '/stripe_webhook'
    ]:
        return
    
    user_id = session.get('user_id')

    if user_id is None:
        g.user = None
    else:
        g.user = db.session.execute(
            db.select(User).filter_by(id=user_id)
        ).scalar()

This change lets Flask instantly serve static files without making any database queries. The g object used in this function is a special object that Flask uses to store data that is accessible throughout the request lifecycle. This object is cleared after each request, so it is safe to store data that is only needed for the current request.

Further optimizations can be made to this function by only loading user data on select routes. However, the username is displayed in the navbar, which is present on every page. We could use session cookies to store the user’s username, but this cookie would need to be evicted when the username changes. This would require a database query or some form of synchronization between clients to check if the username has changed, which would further complicate the logic and negate the performance benefits of using a cookie.

Database Optimizations

After testing, I noticed that one of the slowest parts of the app was loading blog posts. According to the profiler, converting the blog post from Markdown to HTML often took several hundred milliseconds. This was because the Markdown was being converted to HTML on every request, even if the blog post had not changed.

To avoid this, I added a column to the blog post table to store the HTML version of the blog post. This column is updated whenever the blog post is created or updated. This change significantly reduced the time taken to load blog posts, with cache hits taking less than 40ms.

class Post(db.Model):
    __tablename__ = 'post'

    id = Column(Integer, primary_key=True, autoincrement=True)
    user_id = Column(Integer, ForeignKey('users.id', ondelete='CASCADE'), nullable=False)
    created = Column(DateTime, nullable=False, server_default=func.current_timestamp())
    title = Column(Text, nullable=False)
    body = Column(Text, nullable=False)
    body_html = Column(Text, nullable=False) # new column for HTML version of the post

    user = relationship('User', back_populates='posts')

Another optimization was made when making queries from routes under a user’s subdomain. Although usernames in MarkWeb are case-sensitive, the subdomain is case-insensitive. This means that the subdomain user.markweb.test is the same as USER.markweb.test. To prevent creating usernames that would result in duplicate subdomains and to speed up queries, each username is stored in lowercase alongside the cased username.

class User(db.Model):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True, autoincrement=True)
    username = Column(Text, unique=True, nullable=False)
    username_lower = Column(Text, unique=True, nullable=False)
    email = Column(Text, unique=True, nullable=False)
    password = Column(Text, nullable=False)
    
    # ...
    
    def __init__(self, *args, **kwargs):
        super(User, self).__init__(*args, **kwargs)
        self.username_lower = self.username.lower()

By defining the constructor method, we can ensure that the username_lower field is always set to the lowercase version of the username. This way, we can query the database using the lowercase version of the username, which is faster and more efficient than using a case-insensitive query. After all optimizations, blogs load much faster, with DOMContentLoaded times reduced to under 500ms.

Database indices were also added to speed up queries. For example, the username_lower field in the User table was indexed to speed up queries for user subdomains.

Usability

Only a small change was made to improve usability in this update, but the change is similar to more to come.

Now, users can upload markdown files directly to the blog post editor. This feature is useful for users who prefer to write blog posts in their favorite text editor. I plan to add more ways to import content into MarkWeb, such as importing from cloud storage, GitHub, zipped directories, and other blogging platforms. The same will be true for options to export content from MarkWeb, such as exporting to PDF, Word, and other formats. The goal is to ensure user ownership of their content first.

The below code adds an input element to the blog post editor that allows users to upload a markdown file. When a file is uploaded, the contents of the file are read and set as the value of the editor, without sending the file to the server.

<div>
    <label for="md_file">Upload Markdown File</label>
    <input type="file" id="md_file" accept=".md,.markdown,.txt" style="display: none;">
</div>
<script>
    document.getElementById('md_file').addEventListener('change', function(event) {
        const file = event.target.files[0];
        if (file) {
            const reader = new FileReader();
            reader.onload = function(e) {
                simplemde.value(e.target.result);
                event.target.value = '';
            };
            reader.readAsText(file);
            
        }
    });
</script>

Conclusion

With the backend in a more stable state, I plan to focus on the frontend in the next update. This will include improving the import/export features, adding more customization options, and adding analytics.

With more time to work on MarkWeb, I hope to bring it to users in the coming week.

The Stack

• Backend: Python/Flask
• Frontend: HTML, CSS, JS
• Database: PostgreSQL
• Hosting: Azure (already working in VS Code and more familiar with Azure)

Here’s a quick architecture rundown:

My Flask app is on a private GitHub repo, and directly deploys to my Azure web app with GitHub Actions.

I have a PostgreSQL server also running on Azure, which I connect to from the Flask app with SQLAlchemy.

I also had an communication/email service set up on Azure, which would provide me with a free 100 outbound emails a day from my custom domain, but it does not include an inbox and is difficult to scale (requires filing customer support tickets with Microsoft). I got a Google Workspace account for this, which provides limits that are plenty high for my current stage. I’m currently sending emails with Flask-Mail through SMTP, but will need to switch to Gmail’s API as Google is phasing out SMTP support for Google Workspace users.

Later, I plan to store user-provided images in Azure Blob Storage and potentially add other Azure networking services.

Dynamic Subdomains

Like GitHub Pages, I planned to give users their own subdomain at <username>.markweb.app, which would give users a feeling of ownership over their own blogs and content. I kicked implementing this further and further down the road, until a couple of days ago. Implementing this seemed really simple, as it only requires a minor change in the backend code:

# Subdirectory routing
@bp.route('/<username>')
def profile(username):
    # ...
    
# Subdomain routing!
@bp.route('/', subdomain='<username>')
def profile(username):
    # ...

This works just fine when locally hosting, with some additional changes in app configuration. First to support subdomains, we need to point some URL at localhost:5000. This is because 127.0.0.1:5000 and localhost:5000 aren’t proper domains, and browsers won’t resolve subdomain routing from them (correction: as I write this, I realize that *.localhost:5000 does resolve in Chrome, but I go through this process anyways to make sure url_for generates proper URLs and SERVER_NAME is working in production). This can be done by editing the hosts file (on Windows: %SystemRoot%\System32\drivers\etc\hosts):

127.0.0.1    markweb.test
127.0.0.1    admin.markweb.test
127.0.0.1    dmicz.markweb.test
127.0.0.1    www.markweb.test

These lines help the OS map hostnames to IP addresses. Unfortunately, the hosts file does not support wildcard subdomain entries (tools like dnsmasq or Acrylic provide support for this), which means I need to manually enter subdomains I want to visit. Then, when testing locally, my Flask app is available at markweb.test:5000. Additionally, SERVER_NAME and SESSION_COOKIE_DOMAIN needs to be configured to help Flask find where to point links to and store cookies at:

    SERVER_NAME=os.environ.get('SERVER_NAME', 'markweb.test:5000'),
    SESSION_COOKIE_DOMAIN=os.environ.get('SESSION_COOKIE_DOMAIN', '.markweb.test:5000'),

Session cookies should be stored at .markweb.app in production, which allows sessions to be accessible from all subdomains of markweb.app. Otherwise, each subdomain would store sessions separately.

DNS + SSL Issues

Rather than locking myself in to Vercel or Netlify, I settled for using Azure to host my app and Cloudflare to configure my domain and DNS zone. I briefly considered self-hosting this service as an educational experience, but am pushing that off for later. Although this means I have to waddle through the mess that cloud platforms like Azure can be, I quickly familiarized myself with Microsoft’s documentation and got my app running.

This was fine until I needed to implement dynamic subdomain routing. Adding a custom domain to an Azure web app does not include subdomains, so we need to use hostname *.markweb.app. The * is used for wildcard DNS records, which match to any existing subdomain name. However, Azure has somewhat unclear error messaging when entering this domain that suggests an SSL certificate should be added later.

Microsoft’s documentation confirms that Azure does not provide wildcard SSL certificates, so unfortunately I need to look for one. Just checking with Azure:

$300/yr for a certificate was not happening for me. I originally looked at Azure Front Door to see if it would provide the certificates, but starting at \$35/month ($420/year), it was also a nonstarter.

Thankfully, I figured out that Cloudflare can provide wildcard SSL certificates for free by enabling proxying. Cloudflare also provides extremely helpful documentation on their Cloudflare SSL/TLS. Originally, I avoided setting DNS records to be proxied through Cloudflare as I was verifying my custom domain through Azure, because Azure would struggle with verifying the records. This is because the proxied records would point to Cloudflare’s servers, which forwards traffic to Azure’s. But when switching back to proxied status after setting up the custom domain and setting SSL/TLS mode to Full (Strict), everything worked perfectly.

Because the traffic is not proxied through Cloudflare, Cloudflare provides free edge certificates for between itself and the browser, as well as free origin certificates for its connection with Azure’s servers. These certificates are in text format, so I use openssl to get a .pfx to upload to Azure.

openssl pkcs12 -export -out certificate.pfx -inkey privateKey.key -in certificate.crt

After uploading, everything worked smoothly, for as many subdomains as I needed.

Another option I considered was to use Let’s Encrypt with Certbot, but this would require a VM to run the Certbot client and would be more difficult to automate.

CORS Issues

Another issue comes with using dyanmic subdomains. When I tried to access certain resources from a subdomain, I got a CORS error. Specifically, https://markweb.app/static/site.webmanifest is not accessible from https://<username>.markweb.app/ due to the cross origin request. This is because the browser blocks requests from different origins by default. CORS is important because it prevents malicious websites from accessing your resources, but can be a pain to deal with when you’re trying to access your own resources. To fix this, I added the following to my Flask app:

# at top of app file
from flask_cors import CORS
cors = CORS()
...
# in app config
SERVER_NAME = os.environ.get('SERVER_NAME', 'markweb.test:5000')
...
# after app is created
cors.init_app(app, resources={r"/static/*": {"origins": [f"https://*.{app.config['SERVER_NAME']}", f"http://*.{app.config['SERVER_NAME']}"]}})

In this case, SERVER_NAME was set to markweb.app when hosting, and markweb.test:5000 for local development.

In theory, this allows requests from any subdomain of markweb.app (*.markweb.app) to access the /static/* resources. This could be a minor security risk, as it allows any subdomain to access these resources, but the resources are public anyways and there are XSS protections in place. However, this code does not work in the first place because CORS does not support wildcards in the Access-Control-Allow-Origin header. Instead, I could set the header to Access-Control-Allow-Origin: * to allow all origins to access the resources (potential risk) or set the header to Access-Control-Allow-Origin: https://<username>.markweb.app to allow only the specific subdomain to access the resources per request.

However, as it turns out this issue has been solved in flask-cors back in 2014, and was easier to find in the repo’s issues than in the documentation. The solution is to use regular expressions in the origins list to match the subdomains:

cors.init_app(app, resources={r"/static/*": {"origins": rf".*\.{app.config['SERVER_NAME']}"}})

This regex matches any subdomain of markweb.app, and allows requests from those subdomains to access the /static/* resources.

What’s Next?

If you visit MarkWeb right now, you’ll see that everything is locked down for now. I’m still working on ensuring the safety and integrity of the app before launching, which should be complete within the next week. Stay tuned for more development + engineering updates on this blog, and on my Twitter. Features I’ll be working on this week:

• Account management
• Blog management/stylesheets
• Blog homepage
• RSS feed
• Email notifications

As a sidenote, Flask is a very impressive web framework and makes building web apps both easy and fun. I’m excited to continue working with it and building out MarkWeb.

In this post, I explain the vision transformer (ViT) architecture, which has found its way into computer vision as a powerful alternative to Convolutional Neural Networks (CNNs).

This implementation will focus on classifying the CIFAR-10 dataset, but is adaptable to many tasks, including semantic segmentation, instance segmentation, and image generation. As we will see, training small ViT models is difficult, and the notebook on fine-tuning (later in this post) explains how to get around these issues.

We begin by downloading the CIFAR-10 dataset, and transforming the data to torch.Tensors.

import numpy as np
import torch
import torch.nn as nn
from torchvision import datasets, transforms

import math

transform = transforms.Compose([transforms.ToTensor()])

train_data = datasets.CIFAR10(root='./data/cifar-10', train=True, download=True, transform=transform)
test_data = datasets.CIFAR10(root='./data/cifar-10', train=False, download=True, transform=transform)

train_data, test_data

Files already downloaded and verified
Files already downloaded and verified





(Dataset CIFAR10
     Number of datapoints: 50000
     Root location: ./data/cifar-10
     Split: Train
     StandardTransform
 Transform: Compose(
                ToTensor()
            ),
 Dataset CIFAR10
     Number of datapoints: 10000
     Root location: ./data/cifar-10
     Split: Test
     StandardTransform
 Transform: Compose(
                ToTensor()
            ))

The images are represented as 3 channel (RGB) 32x32 pixel images. The dataset can be indexed, with the first index being the image index, and the second index indexing either the image data or the target. The pixel values are represented as torch.float32 values from 0 to 1.

train_data.data.shape, len(train_data.targets)

((50000, 32, 32, 3), 50000)

train_data[0][0].numpy().shape, train_data[0][1]

((3, 32, 32), 6)

train_data[0][0][0], train_data[0][0].dtype

(tensor([[0.2314, 0.1686, 0.1961,  ..., 0.6196, 0.5961, 0.5804],
         [0.0627, 0.0000, 0.0706,  ..., 0.4824, 0.4667, 0.4784],
         [0.0980, 0.0627, 0.1922,  ..., 0.4627, 0.4706, 0.4275],
         ...,
         [0.8157, 0.7882, 0.7765,  ..., 0.6275, 0.2196, 0.2078],
         [0.7059, 0.6784, 0.7294,  ..., 0.7216, 0.3804, 0.3255],
         [0.6941, 0.6588, 0.7020,  ..., 0.8471, 0.5922, 0.4824]]),
 torch.float32)

If you are familiar with the transformer architecture, you likely know that transformers work with vectors to model different modalities. For a text-based modality, this means somehow tokenizing a string of text into characters or larger chunks, and training an embedding table to represent each token as a vector. We hope that tokenization results in semantic units, so that each vector may represent a concept with a specific meaning. As an example, the string “This is a test.” may tokenize as follows:

This is a test.

To adapt the transformer architecture for image tasks, we need to represent image data as a sequence of vectors, similar to how text is tokenized. In the original ViT paper, the authors address this by dividing an image into many patches and flattening them into vectors. With CIFAR-10, an image $x \in \mathbb{R}^{H\times W\times C}$ is turned into several flattened 2D patches of the form $x_p \in \mathbb{R}^{N\times (P^2\cdot C)}$, where $(H,W)$ are the image dimensions (32x32), $C$ is the number of channels (3 for RGB), and $P$ is the patch size. The number of flattened 2D patches is then $N = \frac{HW}{P^2}$. Finally, we project the flattened patches to latent vectors of size $D$, using the linear projection $\mathbf{E} \in \mathbb{R}^{(P^2\cdot C)\times D}$.

patch_size = 4
for i in range(0, 32, patch_size):
    for j in range(0, 32, patch_size):
        patch = train_data[0][0][:, i:i+patch_size, j:j+patch_size]

        # do something with patch...

print(patch.shape)

torch.Size([3, 4, 4])

A naive implementation of turning the images into patches may look like the code above. However, we can accelerate this process by using torch.Tensor.unfold.

image = torch.arange(0.,48).reshape(1, 3, 4, 4) # batch size, channels, width, height
image, image.unfold(2, 2, 2).unfold(3, 2, 2) # first unfold width, then height into 2x2 patches

(tensor([[[[ 0.,  1.,  2.,  3.],
           [ 4.,  5.,  6.,  7.],
           [ 8.,  9., 10., 11.],
           [12., 13., 14., 15.]],
 
          [[16., 17., 18., 19.],
           [20., 21., 22., 23.],
           [24., 25., 26., 27.],
           [28., 29., 30., 31.]],
 
          [[32., 33., 34., 35.],
           [36., 37., 38., 39.],
           [40., 41., 42., 43.],
           [44., 45., 46., 47.]]]]),
 tensor([[[[[[ 0.,  1.],
             [ 4.,  5.]],
 
            [[ 2.,  3.],
             [ 6.,  7.]]],
 
 
           [[[ 8.,  9.],
             [12., 13.]],
 
            [[10., 11.],
             [14., 15.]]]],
 
 
 
          [[[[16., 17.],
             [20., 21.]],
 
            [[18., 19.],
             [22., 23.]]],
 
 
           [[[24., 25.],
             [28., 29.]],
 
            [[26., 27.],
             [30., 31.]]]],
 
 
 
          [[[[32., 33.],
             [36., 37.]],
 
            [[34., 35.],
             [38., 39.]]],
 
 
           [[[40., 41.],
             [44., 45.]],
 
            [[42., 43.],
             [46., 47.]]]]]]))

We can then reshape this tensor into flat patches seperated by channel. We want to combine patches of the same location by channel, so we then permute the dimensions and reshape once again.

# N = H x W / P^2, or the number of patches
image.unfold(2,2,2).unfold(3,2,2).reshape(1, -1, 4, 4) # B x C x N x (P^2)

tensor([[[[ 0.,  1.,  4.,  5.],
          [ 2.,  3.,  6.,  7.],
          [ 8.,  9., 12., 13.],
          [10., 11., 14., 15.]],

         [[16., 17., 20., 21.],
          [18., 19., 22., 23.],
          [24., 25., 28., 29.],
          [26., 27., 30., 31.]],

         [[32., 33., 36., 37.],
          [34., 35., 38., 39.],
          [40., 41., 44., 45.],
          [42., 43., 46., 47.]]]])

image.unfold(2,2,2).unfold(3,2,2).reshape(1, -1, 4, 4).permute(0, 2, 1, 3) # B x N x C x P
image.unfold(2,2,2).unfold(3,2,2).reshape(1, -1, 4, 4).permute(0, 2, 1, 3).reshape(1, 4, -1) # B x N x (C*P^2)

tensor([[[ 0.,  1.,  4.,  5., 16., 17., 20., 21., 32., 33., 36., 37.],
         [ 2.,  3.,  6.,  7., 18., 19., 22., 23., 34., 35., 38., 39.],
         [ 8.,  9., 12., 13., 24., 25., 28., 29., 40., 41., 44., 45.],
         [10., 11., 14., 15., 26., 27., 30., 31., 42., 43., 46., 47.]]])

Another embedding that is critical for the transformer architecture to understand context is the positional embedding. In the text modality, the positional embedding is often implemented using cosine functions. To represent 2D position however, we will use a standard embedding table that is learned during training. There are other possibilities for representing position, including 2D-aware positional embeddings, but these are harder to implement and result in negligble performance differences.

Finally, we also want to add a specific vector for information about the class of the image to the transformer input. Through each transformer block, this vector is modified, before it is fed into a multi-layer perceptron at the last block to determine the class of the image. We will use a learnable embedding for this and prepend it to the other embedded patches.

class PatchEmbedding(nn.Module):
    def __init__(self, img_size = 32, patch_size = 4, in_chans = 3, embed_dim = 768):
        super().__init__()
        self.img_size   = img_size
        self.patch_size = patch_size    # P
        self.in_chans   = in_chans      # C
        self.embed_dim  = embed_dim     # D

        self.num_patches = (img_size // patch_size) ** 2        # N = H*W/P^2
        self.flatten_dim = patch_size * patch_size * in_chans   # P^2*C
        
        self.proj = nn.Linear(self.flatten_dim, embed_dim) # (P^2*C,D)

        self.position_embed = nn.Parameter(torch.zeros(1, 1 + self.num_patches, embed_dim))
        self.class_embed    = nn.Parameter(torch.zeros(1, 1, embed_dim))

    def forward(self, x):
        B, C, H, W = x.shape

        x = x.unfold(2, self.patch_size, self.patch_size).unfold(3, self.patch_size, self.patch_size)
        x = x.reshape(1, -1, self.patch_size, self.patch_size)
        x = x.permute(0, 2, 1, 3).reshape(B, self.num_patches, -1)

        x = self.proj(x)

        cls_emb = self.class_embed.expand(B, -1, -1)
        x = torch.cat((cls_emb, x), dim = 1)

        x = x + self.position_embed
        return x

patch_embed = PatchEmbedding()

embeddings = patch_embed(torch.stack([train_data[i][0] for i in range(10)]))

embeddings, embeddings.shape

(tensor([[[ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  0.0000e+00,
            0.0000e+00,  0.0000e+00],
          [-3.9495e-01,  2.6499e-01, -9.8315e-02,  ...,  1.9426e-01,
            3.5573e-01, -1.2782e-01],
          [-4.9758e-01,  2.3686e-01, -1.7172e-01,  ...,  3.5498e-02,
            1.5119e-01,  6.9413e-03],
          ...,
          [-2.0157e-01,  1.8175e-01, -1.1618e-01,  ...,  6.3366e-02,
            1.4141e-01, -2.3075e-01],
          [-1.4961e-01,  2.5366e-01, -4.7240e-02,  ...,  3.0900e-02,
            1.3584e-01, -1.5386e-01],
          [-1.2538e-01,  1.3091e-01, -1.2969e-01,  ...,  1.0788e-01,
            5.3261e-02, -1.0373e-01]],
 
         [[ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  0.0000e+00,
            0.0000e+00,  0.0000e+00],
          [-3.2531e-01,  2.7973e-01, -2.7797e-01,  ...,  5.6054e-02,
            2.4562e-01, -5.4279e-02],
          [-3.7752e-01,  3.7624e-01, -1.9639e-01,  ..., -4.6192e-02,
            2.7113e-01, -6.9035e-02],
          ...,
          [-6.4879e-02, -9.6972e-04, -2.6319e-01,  ...,  2.6897e-01,
           -6.9230e-02, -7.6847e-02],
          [-2.5097e-01,  1.2145e-01, -3.3001e-01,  ...,  1.3163e-01,
            3.4776e-01, -5.5025e-02],
          [-2.2934e-01,  1.7568e-01, -2.5048e-01,  ...,  1.0340e-01,
            9.4081e-02, -3.7091e-02]],
 
         [[ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  0.0000e+00,
            0.0000e+00,  0.0000e+00],
          [-4.5876e-01,  2.7328e-01, -1.4612e-01,  ...,  1.1323e-01,
            2.5706e-01, -6.0638e-02],
          [-3.0678e-01,  3.2576e-01, -1.7295e-01,  ..., -2.9040e-02,
            2.4391e-01, -1.0266e-01],
          ...,
          [-5.1798e-01,  1.2231e-01, -2.8323e-01,  ...,  1.2910e-01,
            1.0092e-01, -5.8057e-02],
          [-4.3961e-01,  3.0831e-01, -6.3392e-02,  ...,  1.3885e-01,
            3.7307e-01, -3.5249e-01],
          [-1.5562e-01,  1.1467e-01, -2.3594e-01,  ...,  2.0727e-01,
            1.0254e-01, -2.6508e-02]],
 
         ...,
 
         [[ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  0.0000e+00,
            0.0000e+00,  0.0000e+00],
          [-4.2991e-01,  2.3180e-01, -1.6275e-01,  ...,  1.3453e-01,
            2.4938e-01, -5.4029e-02],
          [-2.7091e-01,  2.7341e-01, -2.3737e-01,  ..., -7.2334e-02,
            1.3330e-01, -1.0379e-01],
          ...,
          [-5.8417e-01,  2.4607e-01, -2.9059e-01,  ...,  1.0830e-02,
            1.0258e-01,  2.7538e-02],
          [-3.6148e-01,  1.1907e-01, -3.6477e-02,  ...,  5.2896e-02,
            4.4220e-02, -9.9179e-02],
          [-4.7744e-01,  1.0193e-01, -1.2614e-01,  ...,  2.0482e-01,
            5.7551e-02,  1.5926e-01]],
 
         [[ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  0.0000e+00,
            0.0000e+00,  0.0000e+00],
          [-9.7016e-01,  4.5935e-01, -2.7754e-01,  ...,  1.2938e-01,
            5.4929e-01, -1.2435e-01],
          [-5.6197e-01,  4.6641e-01, -1.0930e-02,  ...,  1.4106e-01,
            6.7737e-01, -1.1047e-01],
          ...,
          [-6.8178e-03,  1.5677e-01, -1.4674e-01,  ...,  9.6236e-02,
           -1.9692e-02, -1.9123e-01],
          [-9.7842e-02,  1.4272e-01, -3.2074e-01,  ..., -2.7982e-02,
            4.1490e-02, -1.2003e-01],
          [ 7.7636e-02,  7.0059e-02, -9.6226e-02,  ...,  1.1831e-01,
           -1.0042e-02, -1.0293e-01]],
 
         [[ 0.0000e+00,  0.0000e+00,  0.0000e+00,  ...,  0.0000e+00,
            0.0000e+00,  0.0000e+00],
          [-3.6146e-01,  1.9650e-01, -8.9722e-02,  ...,  2.0228e-01,
            2.3418e-01, -2.1788e-02],
          [-5.1785e-01,  7.8911e-02, -2.4814e-01,  ...,  6.2777e-02,
            1.8336e-01,  1.5384e-02],
          ...,
          [ 9.7187e-02,  6.5835e-02, -1.2366e-01,  ...,  8.6316e-02,
           -1.8488e-02, -8.9563e-02],
          [-1.8432e-02,  1.1963e-01, -1.7618e-01,  ..., -3.2793e-02,
           -2.6466e-02, -6.3163e-02],
          [-1.1296e-01,  1.0668e-01, -8.1302e-02,  ...,  1.3928e-01,
            7.5195e-02, -1.1134e-01]]], grad_fn=<AddBackward0>),
 torch.Size([10, 65, 768]))

As seen above, we are able to embed batches into our desired embedding dimension, with a correct number of vectors $N+1$.

We can now continue to implement the standard transformer architecture, with one notable change from GPT-like architectures. The attention mechanism of GPT involves multi-headed causal self-attention, which means that vectors are only allowed to query and interact with previous vectors. Although this makes sense in a language model that wants to extract causal contextual information, we want all vectors to communicate with all other vectors, and want to prevent applying an attention mask. Otherwise, the implementation remains unchanged.

class SelfAttention(nn.Module):
    def __init__(self, embed_dim = 768, num_heads = 4, bias = False, dropout=0.1):
        super().__init__()
        assert embed_dim % num_heads == 0

        self.embed_dim   = embed_dim
        self.num_heads   = num_heads
        self.head_dim    = embed_dim // num_heads

        self.query   = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.key     = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.value   = nn.Linear(embed_dim, embed_dim, bias=bias)

        self.out     = nn.Linear(embed_dim, embed_dim, bias=bias)

        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, N, _ = x.size()

        q = self.query(x).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        k = self.key(x).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        v = self.value(x).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)

        # do NOT use causal attention as we are not dealing with sequential data (image patches are unordered)
        attn = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        attn = attn.softmax(dim=-1)

        out = (attn @ v).permute(0, 2, 1, 3).reshape(B, N, self.embed_dim)

        out = self.out(out)

        return out

MSA = SelfAttention()
LN = nn.LayerNorm(embeddings.shape, bias=False)

MSA(LN(embeddings))

tensor([[[-0.5860,  0.3954,  0.2497,  ...,  0.1984,  0.0070,  0.1630],
         [-0.6067,  0.4070,  0.2400,  ...,  0.2375,  0.0075,  0.1689],
         [-0.5888,  0.3900,  0.2408,  ...,  0.2250,  0.0061,  0.1623],
         ...,
         [-0.5961,  0.3959,  0.2409,  ...,  0.2233,  0.0063,  0.1664],
         [-0.5848,  0.3911,  0.2453,  ...,  0.2099,  0.0056,  0.1627],
         [-0.5837,  0.3875,  0.2426,  ...,  0.2105,  0.0060,  0.1622]],

        [[-0.5661,  0.3710,  0.2337,  ...,  0.2130,  0.0105,  0.1403],
         [-0.5794,  0.3729,  0.2260,  ...,  0.2429,  0.0077,  0.1412],
         [-0.5740,  0.3701,  0.2325,  ...,  0.2309,  0.0033,  0.1431],
         ...,
         [-0.5679,  0.3716,  0.2228,  ...,  0.2239,  0.0129,  0.1400],
         [-0.5739,  0.3683,  0.2193,  ...,  0.2312,  0.0209,  0.1412],
         [-0.5778,  0.3683,  0.2226,  ...,  0.2343,  0.0075,  0.1407]],

        [[-0.5962,  0.3924,  0.2470,  ...,  0.2167,  0.0121,  0.1635],
         [-0.6102,  0.3943,  0.2456,  ...,  0.2575,  0.0085,  0.1720],
         [-0.6101,  0.3979,  0.2475,  ...,  0.2511,  0.0041,  0.1728],
         ...,
         [-0.6054,  0.3946,  0.2409,  ...,  0.2755,  0.0149,  0.1746],
         [-0.6069,  0.3830,  0.2432,  ...,  0.2586,  0.0115,  0.1723],
         [-0.5958,  0.3992,  0.2414,  ...,  0.2405,  0.0254,  0.1696]],

        ...,

        [[-0.5766,  0.3838,  0.2403,  ...,  0.2068,  0.0180,  0.1598],
         [-0.5917,  0.3846,  0.2338,  ...,  0.2433,  0.0178,  0.1638],
         [-0.5838,  0.3799,  0.2379,  ...,  0.2283,  0.0100,  0.1623],
         ...,
         [-0.5829,  0.3680,  0.2296,  ...,  0.2442,  0.0090,  0.1594],
         [-0.5840,  0.3761,  0.2376,  ...,  0.2264,  0.0136,  0.1623],
         [-0.5744,  0.3830,  0.2314,  ...,  0.2276,  0.0190,  0.1604]],

        [[-0.5098,  0.3486,  0.2086,  ...,  0.1954,  0.0156,  0.1372],
         [-0.5640,  0.3621,  0.1949,  ...,  0.2696,  0.0105,  0.1501],
         [-0.5355,  0.3669,  0.2019,  ...,  0.2246,  0.0319,  0.1437],
         ...,
         [-0.5099,  0.3471,  0.2052,  ...,  0.2006,  0.0181,  0.1391],
         [-0.5051,  0.3417,  0.2025,  ...,  0.1998,  0.0165,  0.1374],
         [-0.5016,  0.3411,  0.2031,  ...,  0.1934,  0.0188,  0.1351]],

        [[-0.5958,  0.3998,  0.2388,  ...,  0.2080,  0.0119,  0.1450],
         [-0.6008,  0.3943,  0.2309,  ...,  0.2356,  0.0160,  0.1482],
         [-0.6080,  0.3970,  0.2299,  ...,  0.2437,  0.0152,  0.1512],
         ...,
         [-0.5894,  0.3960,  0.2359,  ...,  0.2095,  0.0150,  0.1462],
         [-0.5887,  0.3914,  0.2346,  ...,  0.2109,  0.0092,  0.1452],
         [-0.5935,  0.3959,  0.2324,  ...,  0.2166,  0.0157,  0.1476]]],
       grad_fn=<UnsafeViewBackward0>)

Finally, we want to implement the multi-layer perceptron and combine all our modules into the transformer block.

class MLP(nn.Module):
    def __init__(self, embed_dim = 768, bias = False, dropout = 0.1):
        super().__init__()
        self.c_fc = nn.Linear(embed_dim, embed_dim * 4, bias=bias)
        self.gelu = nn.GELU()
        self.c_proj = nn.Linear(embed_dim * 4, embed_dim, bias=bias)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        x = self.dropout(x)

        return x
    
class Block(nn.Module):

    def __init__(self, embed_dim = 768, bias = False):
        super().__init__()
        self.ln_1 = nn.LayerNorm(embed_dim, bias=bias)
        self.attn = SelfAttention(embed_dim, bias=bias)
        self.ln_2 = nn.LayerNorm(embed_dim, bias=bias)
        self.mlp = MLP(embed_dim, bias=bias)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x
    
    

We will also need to extract information from the class token mentioned earlier to classify the image according to the 10 CIFAR-10 categories. We do this by acting on the class token (a tensor with dimension equal to the embed_dim) with an MLP. The original ViT paper suggests using a hidden layer during pretraining, but finetuning without a hidden layer. For simplicity, we will use a linear layer in the example below.

The final element to consider in our model is the output. Unlike a transformer model like GPT, we would like to produce a probability distribution of the various image classes in CIFAR-10. To achieve this, we use the class vector prepended to the input that was mentioned earlier. After the vector has passed through each transformer block, we can finally take the class vector and pass it through a linear projection once more to get a probability distribution across all ten image classes.

class ViT(nn.Module):

    def __init__(self, embed_dim = 768, num_layers = 4, out_dim = 10, bias = False, dropout = 0.1):
        super().__init__()

        self.transformer = nn.ModuleDict(dict(
            pe = PatchEmbedding(),
            drop = nn.Dropout(dropout),
            h = nn.ModuleList([Block() for _ in range(num_layers)]),
            ln_f = nn.LayerNorm(embed_dim)
        ))
        self.head = nn.Linear(embed_dim, out_dim, bias=False)


        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))

    def get_num_params(self):
        n_params = sum(p.numel() for p in self.parameters())
        return n_params

    def forward(self, x):
        emb = self.transformer.pe(x)
        x = self.transformer.drop(emb)
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)

        class_token = x[:, 0]
        logits = self.head(class_token)
        return logits

vit = ViT()
vit(torch.stack([train_data[i][0] for i in range(10)]))

number of parameters: 28.42M





tensor([[-0.6155,  0.0833,  0.3612,  0.1262,  0.3042,  0.1746, -0.8091, -0.1743,
          0.7834, -0.0631],
        [-1.0184, -0.1414, -0.1417, -0.2963,  0.3737,  0.0205, -0.6493, -0.0109,
          0.8699, -0.1459],
        [-0.9966,  0.1542, -0.1182, -0.0219,  0.5765,  0.0459, -0.4171, -0.1138,
          0.9797, -0.0116],
        [-0.7114,  0.0771,  0.1259,  0.0709,  0.4262,  0.2157, -0.9031, -0.0468,
          0.7561, -0.0685],
        [-0.7684,  0.0140, -0.1220, -0.2881,  0.6543, -0.0601, -0.4210, -0.1563,
          0.8449,  0.0266],
        [-0.6628,  0.2121, -0.0296,  0.1316,  0.4504,  0.5865, -1.1234, -0.1563,
          0.6179,  0.0671],
        [-0.8797,  0.1741,  0.2067, -0.1092,  0.4820,  0.0855, -0.3633,  0.1939,
          0.8685, -0.4827],
        [-0.8408,  0.1524, -0.1209, -0.2138,  0.3352,  0.1860, -0.5956,  0.1675,
          0.9737,  0.0405],
        [-0.8931, -0.3047, -0.1372, -0.3863,  0.4625,  0.2419, -0.4026,  0.1199,
          0.7785,  0.1355],
        [-0.4177,  0.0746,  0.1027,  0.1930,  0.4062,  0.0012, -0.6935, -0.1216,
          1.0843, -0.0408]], grad_fn=<MmBackward0>)

Finetuning larger ViTs on small datasets

Small ViT models struggle to learn the inductive biases necessary for image recognition on small datasets, such as CIFAR-10, which makes training them difficult. Instead, ViTs are often pretrained on very large image datasets before being finetuned to downstream tasks, similar to GPT-style LLMs. As an example, the original ViT paper uses ImageNet-21k with 21k classes and 14M images as one of its datasets.

More technically, fine-tuning a ViT model typically involves turning the classification head into an MLP with no hidden layer ($D \times K$, where $D$ is the embedding dimension and $K$ is the number of downstream classes). Additionally, position embeddings need to be interpolated in 2D to adjust to different image resolutions.

Here, we will fine-tune the smallest model from the original ViT paper, ViT-B/16, trained on ImageNet-21k. We begin by downloading the CIFAR-10 dataset, scaling images up from 32x32 to 224x224. Typically, images are scaled up beyond the pre-training size at fine-tuning to improve performance, but here we will just keep it at the input dimensions for simplicity.

import torch
from torch import nn
from torch.utils.data import DataLoader

import torchvision
from torchvision import transforms


transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_set = torchvision.datasets.CIFAR10(root='./data/cifar-10-finetune', train=True, download=True, transform=transform)
test_set = torchvision.datasets.CIFAR10(root='./data/cifar-10-finetune', train=False, download=True, transform=transform)

train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=4, pin_memory=True)
test_loader = DataLoader(test_set, batch_size=128, shuffle=False, num_workers=4, pin_memory=True)

We then load the ViT-B/16 model from the torchvision module and load it onto the GPU. We need to set the head of the model to fit our fine-tuning task, which in this case has 10 classes. A simple linear layer is most commonly used, and we make sure to set the first dimension to the final dimension of the previous layer. Then, we use an Adam optimizer with some learning rate to finetune. The original ViT paper discusses the learning rate schedules for fine-tuning in further depth, but a constant rate will suffice for demonstration.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = torchvision.models.vit_b_16(pretrained=True)
model = model.to(device)

model.heads = nn.Linear(model.heads[0].in_features, 10)
model.heads = model.heads.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.003)

To train and evaluate the model, we use simple and generic PyTorch training code.

def train(model, criterion, optimizer, loader, epochs=10):
    model.train()
    for epoch in range(epochs):
        for images, labels in loader:
            images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True)
            outputs = model(images)
            loss = criterion(outputs, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f"Epoch {epoch+1}, Loss: {loss.item()}")

def evaluate(model, loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    return accuracy

At fine-tuning, very little training is actually required, and we can get decent performance from just 1 epoch.

train(model, criterion, optimizer, train_loader, epochs=1)
accuracy = evaluate(model, test_loader)
print(f"Test Accuracy: {accuracy}%")

With this simple fine-tuning approach, we achieve over 95% accuracy on CIFAR-10, which only became SOTA in 2014 with fractional max-pooling. Note that this finetuning was done with all the layers locked (requires_grad=False), except for the classification head. More advanced fine-tuning, involving larger pre-trained models, learning rate schedules, and image resolution adjustments all bring this to what has been SOTA performance since 2020, 99.5% top-1 accuracy. Although this comes at the cost of having to train a huge model and needing extra training data, the DeiT vision transformer models introduced in Training data-efficient image transformers & distillation through attention are much smaller than ViT-H/16, can be distilled from Convnets, and achieve up to 99.1% accuracy on CIFAR-10.

The ability to detect AI-generated text is an important issue, not only because of academic integrity issues, but also due to misinformation, security, and copyright concerns. A new method for detection of machine-generated text, called Binoculars, achieves over 90% accuracy in detection at a 0.01% false positive rate. In this notebook, I annotate key parts of the paper, explaining the mechanisms behind this new method and implementing it piece-by-piece. Code from the original paper is available here and this Jupyter Notebook is available here.

The authors of the paper have also created a Hugging Face Space to try out the method.

Table of Contents

• LLM Detection
• Part 1: Understanding Binoculars
  • Perplexity
  • Motivation
• Part 2: Implementation
  • Prelims
  • Observer and Performer Models
  • Perplexity and cross-perplexity
  • Binoculars Score
• Part 3: Further Investigation
• Conclusion

LLM Detection

The motivation behind LLM Detection is harm reduction, to trace text origins, block spam, and identify fake news produced by LLMs. Preemptive detection methods attempt to “watermark” generated text, but requires full control of the generating models, which already seems to be impossible. Therefore, more recent works have been on post-hoc detection methods, which could be used without the cooperation of the text’s author. The paper’s authors suggest that there are two main groups for post-hoc detectors, the first being finetuning a pretrained language model to perform binary classification. There are many additional techniques that make this approach more effective, but all implementations will require training on text produced by the target model, which is both computationally expensive and limited by the number of new models that are being open-sourced.

The second group uses statistical signatures of machine-generated text, with the aim of zero-shot learning. This would allow for the detection of a wide range of models, with little to no training data. These methods use measures such as perplexity, perplexity curvature, log rank, intrinsic dimensionality, and n-gram analysis. The Binoculars paper proposes a focus on low false positive rate (FPR) and high performance on out-of-domain samples, rather than focusing on classifier AUCs for the high-stakes application of LLM detection.

Part 1: Understanding Binoculars

Perplexity

LLMs parse a string of characters $s$ into tokens using a tokenizer $T$, providing a list of tokens $\vec{x}$. Given this input of tokens, a language model $\mathcal{M}$ predicts the token that is next in the sequence by outputting a probability distribution over its vocabulary, $V$:

A natural way to then predict whether or not text was generated by a machine is to measure how likely it is that the language model $\mathcal{M}$ generates each token in $s$, given all the tokens before it. Perplexity is a common baseline that measures this, and we define log-perplexity ($\log \text{PPL}$) as the average log-likelihood of all tokens in $s$:

Perplexity is a reasonable first approach for a couple of reasons:

• Humans often produce higher perplexity text than LLMs (“more surprising”).
• The loss function LLMs are trained on is the log perplexity, as the models are effectively asked to reproduce its test data, given only a prefix of it.

Perplexity, however, is insufficient alone as an approach to detect LLMs. There are many prompts that may produce high perplexity output, such as those that introduce highly specialized domain knowledge, complex vocabulary, new ideas, or prompts that create output that is highly dependent on the context. Tuning the perplexity threshold high enough to catch these LLM outputs inevitably raises the false positive rate for output that is produced by humans, but is more predictable and less domain specific.

As an example from the original paper, prompts such as “1, 2, 3,” have the lowest perplexity with completions such as “4, 5, 6.” However, a prompt about a capybara that is also an astrophysicist produces surprising output, which has a high perplexity. With context, the perplexity of the output would be much lower, but LLM detection in practice must rely on not having access to the language model’s context or prompt.

Motivation

Binoculars uses a mechanism motivated by the aforementioned problems to estimate “baseline perplexity” that was induced by the context and prompt of a text. By comparing actual perplexity to the expected perplexity, we have a much better idea of whether or not an LLM produced the text in question. This is because we can expect the perplexity of human-written text to be even higher perplexity than that of a machine, given the same prompt and context.

To measure the baseline perplexity, the authors introduce cross-perplexity, which is cross-entropy measured between two models on the same string $s$:

NB: The measure relies on the models having the same tokenizer $T$.

This score in essence measures how surprised $\mathcal{M}_1$ is by the token predictions of $\mathcal{M}_2$, therefore providing some sense of how much perplexity a machine-generated token may have. Combining this with perplexity, we reach the proposed Binoculars score:

With this mechanism, $\mathcal{M}_1$ acts as the observer model, while $\mathcal{M}_2$ acts as the performer model. The score is then a comparison of how surprising the producer of the string is and how surprising the performer model’s predictions of the string are.

Part 2: Implementation

Prelims

%pip install sentencepiece transformers torch numpy gradio gradio_client

from typing import Union
import numpy as np
import torch
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer

# Change according to hardware
DEVICE_1 = "cuda:0"
DEVICE_2 = "cpu"

Observer and Performer Models

We begin implementation by selecting the two language models, the observer ($\mathcal{M}_1$) and performer ($\mathcal{M}_2$). This implementation will stick to the choices¹ of Falcon-7B-instruct as an observer and Falcon-7B as a performer, which is the highest performing combination in the Binoculars paper. We first verify that both tokenizers are identical.

torch.set_grad_enabled(False)

observer_name = "tiiuae/falcon-7b-instruct"
performer_name = "tiiuae/falcon-7b"

identical_tokens = (AutoTokenizer.from_pretrained(observer_name).vocab ==
                    AutoTokenizer.from_pretrained(performer_name).vocab)

identical_tokens

True

observer_model = AutoModelForCausalLM.from_pretrained(observer_name,
                                                                   device_map={"": DEVICE_1},
                                                                   trust_remote_code=True,
                                                                   torch_dtype=torch.bfloat16)

performer_model = AutoModelForCausalLM.from_pretrained(performer_name,
                                                                     device_map={"": DEVICE_2},
                                                                     trust_remote_code=True,
                                                                     torch_dtype=torch.bfloat16)

observer_model.eval()
performer_model.eval()

tokenizer = AutoTokenizer.from_pretrained(observer_name)

We then instantiate the tokenizer for both models, which turns the input text into several tokens.

def tokenize(text):
    return tokenizer(text, return_tensors="pt")

tokenize("Hello, my dog is cute")

{'input_ids': tensor([[9856,   23,  491, 3696,  304, 7209]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1]])}

Perplexity and cross-perplexity

To implement perplexity, we need to get the observer model’s log probability for every word in the sequence. Using the observer model, we first obtain the logits for some string.

@torch.inference_mode()
def get_logits(encodings):
    observer_logits = observer_model(**encodings.to(DEVICE_1)).logits
    performer_logits = performer_model(**encodings.to(DEVICE_2)).logits
    return observer_logits, performer_logits

encoding = tokenize('''Dr. Capy Cosmos, a capybara unlike any other, astounded the scientific community with his 
groundbreaking research in astrophysics. With his keen sense of observation and unparalleled ability to interpret 
cosmic data, he uncovered new insights into the mysteries of black holes and the origins of the universe. As he 
peered through telescopes with his large, round eyes, fellow researchers often remarked that it seemed as if the 
stars themselves whispered their secrets directly to him. Dr. Cosmos not only became a beacon of inspiration to 
aspiring scientists but also proved that intellect and innovation can be found in the most unexpected of creatures.'''[:100])


observer_logits, performer_logits = get_logits(encoding)
observer_logits, performer_logits

(tensor([[[-23.1250, -18.1250,  -9.6250,  ..., -10.6875, -12.1250,  -9.1875],
          [-13.7500, -19.6250, -14.1875,  ..., -15.0000, -16.6250, -10.9375],
          [-13.0625, -16.8750, -14.8750,  ..., -17.8750, -15.3750, -14.5000],
          ...,
          [-12.9375, -12.3125, -12.6875,  ..., -16.0000, -14.8750, -18.0000],
          [-16.7500, -15.5000, -16.2500,  ..., -16.7500, -18.3750, -15.8125],
          [-16.1250, -17.6250, -14.5000,  ..., -18.1250, -19.1250, -17.2500]]],
        device='cuda:0', dtype=torch.bfloat16),
 tensor([[[-20.3750, -21.0000, -14.5625,  ..., -13.2500, -15.8750,  -8.5625],
          [-13.0625, -18.7500, -15.1875,  ..., -15.3125, -17.1250, -11.5625],
          [-12.3750, -17.6250, -16.6250,  ..., -17.3750, -16.1250, -14.3125],
          ...,
          [-10.9375, -12.4375, -11.9375,  ..., -14.1875, -14.6875, -16.6250],
          [-14.6250, -13.8750, -16.7500,  ..., -17.6250, -19.1250, -15.6250],
          [-13.6250, -15.5000, -14.5625,  ..., -18.5000, -19.0000, -16.6250]]],
        dtype=torch.bfloat16))

encoding.input_ids.shape, observer_logits.shape

(torch.Size([1, 26]), torch.Size([1, 26, 65024]))

These logits are in the shape of $B \times S \times V$, where $B$ is the batch size, $S$ is the sequence token length, and $V$ is the size of the vocabulary. Because we are only processing a single string, the batch size is 1, and the first 100 characters were tokenized into 26 tokens. We are interested in the logits of each token in the sequence (which provides a prediction of the next token) and we want to $S$ vectors of size $V$ (in this case, ~65k) that contains the log probabilities of each of the tokens in the vocabulary.

S = observer_logits.shape[-2]
V = observer_logits.shape[-1]

observer_logits[..., :-1, :].contiguous().shape

torch.Size([1, 25, 65024])

We then compare these to the true next tokens in the sequence (from the encoding of the string), ignoring the first token because it is not being predicted.

encoding.input_ids[..., 1:].shape

torch.Size([1, 25])

We will ignore the attention mask for now, but it’s a key component later once we use larger batches. To compute the perplexity, we use the equation written earlier:

Thankfully, PyTorch provides a class for this, torch.nn.CrossEntropyLoss. The class can then be called with the model outputs with shape $(S, V)$ and the target output in the shape of $(S)$. We transpose the logit tensor’s last two dimensions to match this shape.

loss = torch.nn.CrossEntropyLoss(reduction='none')

ppl = loss(observer_logits[..., :-1, :].contiguous().transpose(1, 2).to("cpu"), 
     encoding.input_ids[..., 1:].contiguous().to("cpu")).float()

ppl, ppl.sum(1)

(tensor([[ 0.7148,  9.8750,  6.0312,  9.8125,  1.8359,  2.9688,  1.9375,  6.9375,
           0.0430,  0.0270, 11.5000,  0.2676,  0.1396,  0.3066,  9.0625,  1.7656,
           1.1250,  2.4531,  0.7109,  1.0859,  1.4297,  4.7188,  9.6875,  5.5938,
           8.8750]]),
 tensor([98.9049]))

From this, we get a measure of how “surprising” each token is to the observer model when performing next-token prediction and sum it together. While some tokens have a perplexity as low as 0.02, others are above 10.

We now implement cross-perplexity by running the same cross entropy function, except the target tensor is no longer a one-hot vector (the observed next token given by the user), but rather the softmax of the observer’s logits.

softmax = torch.nn.Softmax(dim=-1)

performer_probs = softmax(performer_logits).view(-1, V)

performer_probs, performer_probs.shape

(tensor([[5.8265e-12, 3.1122e-12, 1.9500e-09,  ..., 7.2177e-09, 5.2387e-10,
          7.8604e-07],
         [7.5903e-08, 2.5648e-10, 9.0804e-09,  ..., 7.9744e-09, 1.3024e-09,
          3.4086e-07],
         [6.6124e-08, 3.4743e-10, 9.4587e-10,  ..., 4.4565e-10, 1.5571e-09,
          9.5461e-09],
         ...,
         [1.3039e-07, 2.8987e-08, 4.7730e-08,  ..., 5.0350e-09, 3.0559e-09,
          4.4020e-10],
         [2.1071e-08, 4.4703e-08, 2.5175e-09,  ..., 1.0477e-09, 2.3465e-10,
          7.7416e-09],
         [2.5332e-07, 3.8883e-08, 9.9186e-08,  ..., 1.9354e-09, 1.1714e-09,
          1.2631e-08]], dtype=torch.bfloat16),
 torch.Size([26, 65024]))

observer_scores = observer_logits.view(-1, V).to("cpu")
observer_scores, observer_scores.shape

(tensor([[-23.1250, -18.1250,  -9.6250,  ..., -10.6875, -12.1250,  -9.1875],
         [-13.7500, -19.6250, -14.1875,  ..., -15.0000, -16.6250, -10.9375],
         [-13.0625, -16.8750, -14.8750,  ..., -17.8750, -15.3750, -14.5000],
         ...,
         [-12.9375, -12.3125, -12.6875,  ..., -16.0000, -14.8750, -18.0000],
         [-16.7500, -15.5000, -16.2500,  ..., -16.7500, -18.3750, -15.8125],
         [-16.1250, -17.6250, -14.5000,  ..., -18.1250, -19.1250, -17.2500]],
        dtype=torch.bfloat16),
 torch.Size([26, 65024]))

We use all the log probabilities and observer softmax scores except for the last values, which correspond to the token after the last one.

xppl = loss(observer_scores[:-1], performer_probs[:-1]).view(-1, S - 1)

xppl, xppl.sum(1)

(tensor([[3.1406, 7.3438, 5.1875, 7.5000, 3.9375, 5.0312, 6.2500, 7.1250, 0.2930,
          0.2598, 4.4688, 1.2188, 0.9102, 1.6719, 4.4688, 1.8281, 4.0312, 4.6875,
          1.0078, 2.7656, 1.9531, 6.0000, 3.8125, 9.8125, 0.9219]],
        dtype=torch.bfloat16),
 tensor([95.5000], dtype=torch.bfloat16))

Binoculars Score

To get the Binoculars score from our previous work, we only need to divide the perplexity score from the observer by the cross perplexity score between the observer and performer model.

binocular_score = ppl.sum(1) / xppl.sum(1)

binocular_score

tensor([1.0357])

From this point, the only thing to do is determine a threshold above which to classify text as being machine generated based on the binoculars score.

Let’s package this process into functions to quickly evaluate any given text.

# redefine to handle batch of strings
def tokenize(batch):
    encodings = tokenizer(batch, return_tensors="pt", 
    padding="longest" if len(batch) > 1 else False, truncation=True,
    max_length=512, return_token_type_ids=False).to(DEVICE_1)
    return encodings

# redefinition with cuda sync
@torch.inference_mode()
def get_logits(encodings):
    observer_logits = observer_model(**encodings.to(DEVICE_1)).logits
    performer_logits = performer_model(**encodings.to(DEVICE_2)).logits
    torch.cuda.synchronize()

    return observer_logits, performer_logits

loss_fn = torch.nn.CrossEntropyLoss(reduction='none')
softmax_fn = torch.nn.Softmax(dim=-1)

def perplexity(encoding, logits):
    shifted_logits = logits[..., :-1, :].contiguous()
    shifted_labels = encoding.input_ids[..., 1:].contiguous()
    shifted_attention_mask = encoding.attention_mask[..., 1:].contiguous()

    ppl = loss_fn(shifted_logits.transpose(1, 2).to("cpu"), shifted_labels) * shifted_attention_mask
    ppl = ppl.sum(1) / shifted_attention_mask.sum(1)
    
    return ppl.to("cpu").float().numpy()

def cross_perplexity(observer_logits, performer_logits, encoding):
    V = observer_logits.shape[-1]
    S = observer_logits.shape[-2]

    performer_probs = softmax_fn(performer_logits).view(-1, V).to("cpu")
    observer_scores = observer_logits.view(-1, V).to("cpu")
    
    xppl = loss_fn(observer_scores, performer_probs).view(-1, S)
    padding_mask = (encoding.input_ids != tokenizer.pad_token_id).type(torch.uint8)
    
    xppl = (xppl * padding_mask).sum(1) / padding_mask.sum(1)
    
    return xppl.to("cpu").float().numpy()

def binocular_score(text):
    batch = [text] if isinstance(text, str) else text
    encodings = tokenize(batch)
    observer_logits, performer_logits = get_logits(encodings)
    ppl = perplexity(encodings, observer_logits)
    xppl = cross_perplexity(observer_logits, performer_logits, encodings)

    return (ppl / xppl).tolist()

tokenizer.pad_token = tokenizer.eos_token