Update app.py

app.py

@@ -1,32 +1,40 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import black
 
-# Load the model and tokenizer from Hugging Face Hub
-model_name = "microsoft/CodeGPT-small-py"  # Model name you provided
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+# Load model and tokenizer
+model_name = "microsoft/CodeGPT-small-py"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Initialize the pipeline
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer, temperature=0.5, top_p=0.9, max_length=150)
 
-def generate_text(prompt):
-    # Generate text using the model
-    generated = generator(prompt, max_length=50)
-    return generated[0]['generated_text']
+def generate_code_with_feedback(prompt):
+    generated_code = generator(prompt, num_return_sequences=1)[0]['generated_text']
+    # Apply self-check for code quality
+    formatted_code = format_code(generated_code)
+    return formatted_code
+
+def format_code(code):
+    return black.format_str(code, mode=black.Mode())
 
 # Streamlit UI
-st.title("CodeGPT Text Generation")
-st.write("Enter a prompt to generate code or text:")
+st.title("Smart Code Generation and Fixing")
+st.write("Enter a prompt to generate or fix code:")
 
-# Text input for the user to enter a prompt
-prompt = st.text_area("Prompt", "def example_function():\n    # Your code here...")
+option = st.radio("Select Action", ("Generate Code", "Fix Code"))
+if option == "Generate Code":
+    prompt = st.text_area("Prompt", "Write a Python function that reverses a string:")
+else:
+    prompt = st.text_area("Prompt", "Fix the following buggy Python code:\n\ndef reverse_string(s):\n    return s[::-1]")
 
-# Button to trigger the model inference
 if st.button("Generate Code"):
     if prompt:
-        generated_text = generate_text(prompt)
+        generated_code = generate_code_with_feedback(prompt)
         st.subheader("Generated Code")
-        st.write(generated_text)
+        st.write(generated_code)
     else:
-        st.warning("Please enter a prompt to generate code.")
+        st.warning("Please enter a prompt.")
+