Update app.py

app.py

@@ -1,12 +1,15 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-#to load the model pipeline
+# Load model and tokenizer
 @st.cache_resource()
 def load_model():
-    return pipeline("text2text-generation", model="Salesforce/codet5p-220m")
+    model_name = "Salesforce/codet5-small"  # A better model for code correction
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    return model, tokenizer
 
-model = load_model()
+model, tokenizer = load_model()
 
 # Streamlit UI
 st.title("CodeCorrect AI")
@@ -16,11 +19,13 @@ code_input = st.text_area("Enter your code here:", height=200)
 
 if st.button("Correct Code"):
     if code_input.strip():
-        response = model(code_input, max_length=512)
-        corrected_code = response[0]['generated_text']
+        # Tokenize and generate corrected code
+        inputs = tokenizer(code_input, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        outputs = model.generate(**inputs, max_length=512)
+        corrected_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
         st.text_area("Corrected Code:", corrected_code, height=200)
     else:
         st.warning("Please enter some code.")
 
-
-
+st.markdown("Powered by Hugging Face 🤗")