diff --git a/.gitignore b/.gitignore
index 519c11c..487fff4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -16,3 +16,5 @@ build/
 .idea/
 
 .env
+
+.vscode/
diff --git a/.streamlit/config.toml b/.streamlit/config.toml
index 216cdf0..cad350f 100644
--- a/.streamlit/config.toml
+++ b/.streamlit/config.toml
@@ -1,2 +1,7 @@
 [theme]
-base="light"
\ No newline at end of file
+base="light"
+primaryColor="#F63366"
+backgroundColor="#FFFFFF"
+secondaryBackgroundColor="#F0F2F6"
+textColor="#262730"
+font="sans serif"
\ No newline at end of file
diff --git a/main.py b/main.py
index 5c71e04..0e0faf2 100644
--- a/main.py
+++ b/main.py
@@ -6,15 +6,26 @@
 
 columns = []
 
+
+def get_max_line_height(responses):
+    max_num_lines = 1
+    for response in responses:
+        content = response["content"]
+        num_lines = len(content.split("\n"))
+        max_num_lines = max(max_num_lines, num_lines)
+    return 35 * max_num_lines
+
+
 def grade_drafts(gladiator, responses):
     with st.spinner("Grading Drafts..."):
         columns = st.columns(3)
+        max_line_height = get_max_line_height(responses=responses)
         for i, column in enumerate(columns):
-            draft_template = drafts_template(i, responses[i])
+            draft_template = drafts_template(i, responses[i], max_line_height)
             with column:
                 st.markdown(draft_template, unsafe_allow_html=True)
 
-        grades_json = gladiator.grade_drafts() 
+        grades_json = gladiator.grade_drafts()
         return grades_json, columns
 
 
@@ -23,27 +34,43 @@ def select_winner(grades_json, columns):
 
     for i, column in enumerate(columns):
         is_winner = i == winning_index
-        grade_template = grades_template(grades_json[i]['score'], grades_json[i]['explanation'], is_winner)
+        grade_template = grades_template(
+            grades_json[i]["score"], grades_json[i]["explanation"], is_winner
+        )
         with column:
             st.markdown(grade_template, unsafe_allow_html=True)
 
     st.subheader("Winning Response:")
     with st.container():
-        st.write(winning_content)
+        st.code(winning_content)
+
 
 def generate_drafts(gladiator):
     with st.spinner("Generating Drafts..."):
-        
-        responses = gladiator.generate_drafts(prompt) 
+        responses = gladiator.generate_drafts(prompt)
         return responses
 
 
 dotenv.load_dotenv()
 
 st.set_page_config(layout="wide")
+
 st.markdown(stylesheet, unsafe_allow_html=True)
+
+# Prism Library for Code Formatting
+st.markdown(
+    """
+        <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/prism/1.24.1/themes/prism.min.css">
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/prism/1.24.1/prism.min.js"></script>
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/prism/1.17.1/components/prism-python.min.js"></script>
+        """,
+    unsafe_allow_html=True,
+)
+
 st.title("Gladiator: May the best response win")
-prompt = st.text_area("Enter a prompt and get the best answer", defaults.default_question, height=300)
+prompt = st.text_area(
+    "Enter a prompt and get the best answer", defaults.default_question, height=300
+)
 
 
 if st.button("Generate best answer"):
@@ -55,9 +82,3 @@ def generate_drafts(gladiator):
     except Exception as e:
         print("Error: ", e)
         st.error("An error occurred: {}".format(str(e)))
-
-
-
-
-
-
diff --git a/ui.py b/ui.py
index e353176..9b2b063 100644
--- a/ui.py
+++ b/ui.py
@@ -1,10 +1,12 @@
-def drafts_template(i, content):
-    return f"<div class='draft'><div><b>Draft {i + 1}</b></div>{content['content']}</div>"
+def drafts_template(i, response, max_line_height):
+    return f"<div class='draft' style='height:{max_line_height}px;'><div><b>Draft {i + 1}</b></div><div><pre><code class='language-python'>{response['content']}</code></pre></div></div>"
+
 
 def grades_template(score, explanation, is_winner):
-    class_name = 'winner' if is_winner else ''
+    class_name = "winner" if is_winner else ""
     return f"<div class='{class_name} score'><div><b>Score {score}</b></div>{explanation}</div>"
 
+
 stylesheet = """
 <style>
 .draft {
@@ -25,4 +27,4 @@ def grades_template(score, explanation, is_winner):
     border: 3px solid #33cc33; 
 }
 </style>
-"""
\ No newline at end of file
+"""