First app version

app.py

@@ -0,0 +1,91 @@
+import streamlit as st
+from io import StringIO 
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from torch.nn import functional as F
+import torch
+import numpy as np
+
+import numpyAc
+
+st.set_page_config(layout="wide")
+
+@st.cache_resource
+def load_model():
+    return AutoModelForCausalLM.from_pretrained(
+        "codellama/CodeLlama-7b-Python-hf",
+        device_map='auto',
+    )
+
+@st.cache_resource
+def load_tokenizer():
+    return AutoTokenizer.from_pretrained("codellama/CodeLlama-7b-Python-hf")
+
+model = load_model()
+tokenizer = load_tokenizer()
+
+st.title('Python file compressor')
+encode_col, decode_col = st.columns(2, gap='medium')
+
+@st.cache_data
+def encode(text):
+    codec = numpyAc.arithmeticCoding()
+    tokenized = tokenizer(text, return_tensors='pt').input_ids.to('cuda')
+    output = list()
+    past_key_values = None
+
+    for i in range(tokenized.shape[1]):
+        with torch.no_grad():
+            output_ = model(
+                input_ids=tokenized[:, i:i + 1],
+                use_cache=True,
+                past_key_values=past_key_values
+            )
+        past_key_values = output_.past_key_values
+        logits = output_.logits[0, -1:, :]
+        output.append(logits)
+    output = torch.cat(output, dim=0)
+    output = F.softmax(output, dim=-1)
+    tokenized = torch.cat((tokenized.squeeze()[1:], torch.tensor([2], device='cuda'))) # Add EOS
+    tokenized = tokenized.type(torch.int16).cpu().numpy()
+    byte_stream, _ = codec.encode(output.cpu(), tokenized)
+    return byte_stream
+
+@st.cache_data
+def decode(byte_stream):
+    decodec = numpyAc.arithmeticDeCoding(byte_stream, 32_000)
+    input_ids = [1]
+    past_key_values = None
+
+    while input_ids[-1] != 2:
+        with torch.no_grad():
+            output = model(
+                input_ids=torch.tensor([input_ids[-1:]], device='cuda'),
+                use_cache=True,
+                past_key_values=past_key_values
+            )
+        past_key_values = output.past_key_values
+        logits = output.logits[0, -1:, :]
+        logits = F.softmax(logits, dim=-1).cpu()
+        next_token = decodec.decode(logits)
+        input_ids.append(next_token)
+    return input_ids
+
+with encode_col:
+    st.header('Convert your python file to binary.')
+    python_file = st.file_uploader("Upload your python file here. I recommend files up to 50-100 lines, so it doesn't take too long.")
+    if python_file is not None:
+        stringio = StringIO(python_file.getvalue().decode("utf-8"))
+        code = stringio.read()
+        bytes_stream = encode(code)
+        bin_filename = f'{python_file.name.split(".")[0]}.bin'
+        st.download_button('Download binary file', bytes_stream, bin_filename)
+
+with decode_col:
+    st.header('Convert your binary file to python')
+    binary_file = st.file_uploader('Upload your binary file here')
+    if binary_file is not None:
+        tokens = decode(binary_file.read())
+        decompressed = tokenizer.decode(tokens, skip_special_tokens=True)
+        py_filename = f'{binary_file.name.split(".")[0]}.py'
+        st.download_button('Download python file', decompressed, py_filename)
+        st.code(decompressed)