Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,7 +1,6 @@
 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"
@@ -12,44 +11,19 @@ st.title("Pseudo-code to Code")
 
 prompt = st.text_area("Enter a code:")
 
-
-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)
+# Fixed maximum length
+max_length = 150
 
 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, 
+            **inputs,
+            max_length=max_length,
+            do_sample=True,
             temperature=0.7
         )
         generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # 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")
+        st.code(generated_code, language="python")
     else:
-        st.warning("Enter prompt.")
+        st.warning("Please enter a prompt.")