Update app.py

app.py

@@ -1,26 +1,30 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
+# Load model and tokenizer
 @st.cache_resource()
 def load_model():
-    model_name = "bigcode/starcoder"
+    model_name = "deepseek-ai/deepseek-coder-6.7b-instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
     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 code here:", height=200)
+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)
+        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.")