app.py

import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load model and tokenizer
@st.cache_resource()
def load_model():
    model_name = "deepseek-ai/deepseek-coder-6.7b-instruct"
    tokenizer = AutoTokenizer.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 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 🤗")