Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,23 +1,55 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import re
 
 # Load model and tokenizer
 MODEL_NAME = "Johannes/code-generation-model-fine-tuned-to-produce-good-code-snippets"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-st.title("Code Generation Model")
+st.title("Pseudo-code to Code")
 
-prompt = st.text_area("Enter a code prompt:")
+prompt = st.text_area("Enter a code:")
 
-max_length = st.slider("Maximum generated tokens", min_value=50, max_value=500, value=150)
+
+def remove_comments_and_headers(code_text):
+    """
+    Removes:
+    - Python-style comments starting with #
+    - C/C++ style comments starting with //
+    - License or header blocks
+    """
+    lines = code_text.split("\n")
+    code_lines = []
+
+    for line in lines:
+        stripped = line.strip()
+        # Skip empty lines or comment lines
+        if stripped == "":
+            continue
+        if stripped.startswith("#") or stripped.startswith("//"):
+            continue
+        code_lines.append(line)
+    
+    # Join remaining lines
+    return "\n".join(code_lines)
 
 if st.button("Generate Code"):
     if prompt.strip() != "":
         inputs = tokenizer(prompt, return_tensors="pt")
-        outputs = model.generate(**inputs, max_length=max_length, do_sample=True, temperature=0.7)
+        outputs = model.generate(
+            **inputs, 
+            max_length=max_length, 
+            do_sample=True, 
+            temperature=0.7
+        )
         generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        st.code(generated_code, language="python")
+        
+        # Clean code by removing comments / license headers
+        cleaned_code = remove_comments_and_headers(generated_code)
+        
+        st.subheader("Generated Code:")
+        st.code(cleaned_code, language="python")
     else:
-        st.warning("Please enter a prompt.")
+        st.warning("Enter prompt.")