Hugging Face's logo

Spaces:
adarshzolekar
/
AI-Code-Reviewer
Running on Zero

App Files Community
AI-Code-Reviewer
File size: 2,492 Bytes
5e0c977
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import gradio as gr
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "Qwen/Qwen2.5-Coder-0.5B-Instruct"

# Load tokenizer and model once at startup
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)

# Use GPU when available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = model.to(device)
model.eval()


SYSTEM_PROMPT = (
    "You are a concise code reviewer. Given a code snippet, point out bugs, "
    "edge cases, complexity issues, and style improvements. "
    "Be specific and brief."
)


@spaces.GPU
def review_code(code):
    if not code or not code.strip():
        return "Paste some code first."

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT
        },
        {
            "role": "user",
            "content": f"Review this code:\n\n{code}"
        }
    ]

    # Apply the model's chat template
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=4096
    )

    # Move inputs to the same device as the model
    inputs = {
        key: value.to(device)
        for key, value in inputs.items()
    }

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=300,
            do_sample=False,
            temperature=None,
            top_p=None
        )

    # Only decode newly generated tokens
    generated_ids = output_ids[0][inputs["input_ids"].shape[1]:]

    review = tokenizer.decode(
        generated_ids,
        skip_special_tokens=True
    )

    return review.strip()


with gr.Blocks(title="AI Code Reviewer") as demo:

    gr.Markdown(
        "## AI Code Reviewer\n"
        "Paste a code snippet and get feedback on bugs, edge cases, "
        "complexity, and style."
    )

    code_input = gr.Code(
        label="Your code",
        language="python",
        lines=15
    )

    review_btn = gr.Button(
        "Review",
        variant="primary"
    )

    output = gr.Markdown()

    review_btn.click(
        review_code,
        inputs=code_input,
        outputs=output
    )


if __name__ == "__main__":
    demo.launch()