Update app.py

app.py

@@ -5,28 +5,15 @@ import torch
 # Load model and tokenizer
 @st.cache_resource()
 def load_model():
-    model_name = "deepseek-ai/deepseek-coder-6.7b-instruct"
+    model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"  # Smaller version
+    
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name, 
+        torch_dtype=torch.float16, 
+        device_map="auto"  # Automatically assigns to CPU/GPU if available
+    )
+    
     return model, tokenizer
 
 model, tokenizer = load_model()
-
-# Streamlit UI
-st.title("CodeCorrect AI")
-st.subheader("AI-powered Code Autocorrect Tool")
-
-code_input = st.text_area("Enter your buggy code here:", height=200)
-
-if st.button("Correct Code"):
-    if code_input.strip():
-        prompt = f"### Fix the following code:\n{code_input}\n### Corrected version:\n"
-        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512).to("cuda" if torch.cuda.is_available() else "cpu")
-        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 🤗")