Initial push

.gitignore

@@ -0,0 +1 @@
+__pycache__

README.md

@@ -9,4 +9,40 @@ app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Early Exit Computational Savings Demo
+
+## Overview
+
+This project demonstrates the concept of "early exiting" in deep learning models to save computational resources without significantly compromising on model performance. Early exit strategies allow a neural network to make predictions at intermediate layers for easy-to-classify instances, thus reducing the overall computation time and resources needed for inference. The application is built to run with Streamlit, offering an interactive web interface to explore the functionalities of the early exit model.
+
+## Features
+
+- **BranchyModel:** An implementation of a deep learning model with early exit points. This model architecture is designed to evaluate the performance and computational savings of using early exits.
+- **Utility Functions:** A set of utilities to support the model's operation, including data preprocessing and performance evaluation metrics.
+- **Streamlit Application:** A user-friendly web interface to interact with the model, visualize its performance, and understand the benefits of early exits.
+
+## Getting Started
+
+### Prerequisites
+
+Ensure you have Python 3.x installed on your machine. You can install all the required dependencies via:
+
+```bash
+pip install -r requirements.txt
+```
+
+### Running the Application
+
+To run the Streamlit application, execute the following command from the root of the project:
+
+```bash
+streamlit run app.py
+```
+
+The command will start a local web server and open the application in your default web browser, allowing you to interact with the BranchyModel and explore its features.
+
+## Project Structure
+
+- **app.py**: The main Streamlit application script.
+- **requirements.txt**: Lists all the Python dependencies required by the project.
+- **src/**: Contains the source code for the project.

app.py

@@ -0,0 +1,123 @@
+# Save this as app.py and run with `streamlit run app.py`
+import time
+import streamlit as st
+import torch
+import pandas as pd
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from typer import clear
+from annotated_text import annotated_text
+
+st.title("Multi-Head LLM Demo")
+st.markdown("""This is a demo of a multi-head language model with early exit capabilities. 
+            The model is based on the Phi-2 architecture and model is available here : https://huggingface.co/valcore/Branchy-Phi-2.
+            \nThe model has four heads, each of which can be exited early based on a threshold. The graph show the depth of early exit for each token (the deeper being the faster) and the time taken to generate each token.
+            Early exited tokens are annotated with the depth of early exit (with a float smaller than 1, 1 being the deepest)
+            """)
+
+def annotated_to_normal(text):
+    result = ""
+    for elem in text:
+        if isinstance(elem, tuple):
+            result += elem[0]
+        else:
+            result += elem
+    return result
+
+def generate_next_token():
+    print(f"Generating next token from {st.session_state.messages}")
+    inputs = ""
+    for message in st.session_state.messages:
+        inputs += message["role"] + ": " + annotated_to_normal(message["content"]) + "\n"
+    inputs += "Assistant:"
+    print(f"Inputs: {inputs}")
+    inputs = st.session_state.tokenizer.encode(inputs, return_tensors="pt")
+    for i in range(50):
+        start = time.time()
+        outputs = st.session_state.model(inputs)
+        stop = time.time()
+        next_token_logits = outputs.logits[:, -1, :].squeeze()
+        next_token_probs = torch.softmax(next_token_logits, dim=-1)
+        next_token_id = torch.argmax(next_token_probs, dim=-1)
+        if next_token_id == 50256:
+            break
+        print(inputs.shape, next_token_id.shape)
+        inputs = torch.cat([inputs, next_token_id.unsqueeze(0).unsqueeze(-1)], dim=-1)
+        next_token = st.session_state.tokenizer.decode(next_token_id, return_tensors="pt")
+        time_taken = stop - start
+        branch_locations = st.session_state.model.config.branch_locations
+        print(outputs.head_indices)
+        if outputs.head_indices in branch_locations:
+            print(sorted(branch_locations, reverse=True))
+            early_exit = (branch_locations.index(outputs.head_indices) + 1) / len(branch_locations)
+        else:
+            early_exit = 0
+        # Add data to dataframe
+        new_row = pd.DataFrame({"Time taken (in ms)": [time_taken], "Early exit depth": [early_exit]})
+        st.session_state.data = pd.concat([st.session_state.data, new_row], ignore_index=True)
+        yield next_token, early_exit
+
+@st.cache_resource
+def load_model(model_str, tokenizer_str):
+    model = AutoModelForCausalLM.from_pretrained(model_str, trust_remote_code=True)
+    model.eval()
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_str)
+    return model, tokenizer
+
+model_str = "valcore/Branchy-Phi-2"
+tokenizer_str = "microsoft/Phi-2"
+
+if "model" not in st.session_state or "tokenizer" not in st.session_state:
+    print("Loading model...")
+    st.session_state.model, st.session_state.tokenizer = load_model(model_str, tokenizer_str)
+
+# Initialize chat history and dataframe
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+st.session_state.data = pd.DataFrame(columns=["Time taken (in ms)", "Early exit depth"])
+
+col1, col2 = st.columns([1, 4])
+
+with col1:
+    early_exit = st.checkbox("Early exit", value=False)
+    if early_exit:
+        st.session_state.model.head_thresholds = [2.506962537765503, 2.656052589416504, 1.924393653869629, 1.4434680938720703]
+    else:
+        st.session_state.model.head_thresholds = [10., 10., 10., 10.]
+    clear_session = st.button("Clear session")
+    if clear_session:
+        print("Clearing session")
+        st.session_state.messages = []
+        st.session_state.data = pd.DataFrame(columns=["Time taken (in ms)", "Early exit depth"])
+
+with col2:
+    # Display chat messages from history on app rerun
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            annotated_text(message["content"])
+
+    prompt = st.chat_input("What is up?")
+    # React to user input
+    if prompt:
+        # Display user message in chat message container
+        with st.chat_message("User"):
+            st.markdown(prompt)
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "User", "content": prompt})
+
+        # Display assistant response in chat message container
+        with st.chat_message("Assistant"):
+            response = []
+            with st.spinner('Running inference...'):
+                for next_token, early_exit in generate_next_token():
+                    if early_exit > 0.0:
+                        response.append(tuple((next_token, str(early_exit))))
+                    else:
+                        response.append(next_token)
+                    print(response)
+            annotated_text(response)
+
+        # Add assistant response to chat history
+        st.session_state.messages.append({"role": "Assistant", "content": response})
+        st.line_chart(st.session_state.data, x=None, y=["Time taken (in ms)", "Early exit depth"])
+        print(st.session_state.messages)

requirements.txt

@@ -0,0 +1,5 @@
+streamlit==1.31.0
+torch==2.0.1
+pandas==2.0.3
+transformers==4.36.0
+st-annotated-text