Spaces:
Sleeping
Sleeping
File size: 4,113 Bytes
a77dc20 |
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 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from huggingface_hub import login
import os
import logging
from datetime import datetime
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Environment variables
HF_TOKEN = os.getenv("HUGGING_FACE_TOKEN")
MODEL_NAME = os.getenv("MODEL_NAME", "google/gemma-2b-it")
class CodeReviewer:
def __init__(self):
self.model = None
self.tokenizer = None
self.device = "cpu"
self.initialize_model()
def initialize_model(self):
"""Initialize the model and tokenizer."""
try:
if HF_TOKEN:
login(token=HF_TOKEN)
logger.info("Loading tokenizer...")
self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
logger.info("Loading model...")
self.model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
device_map={"": self.device},
torch_dtype=torch.float32,
low_cpu_mem_usage=True
)
logger.info("Model loaded successfully")
except Exception as e:
logger.error(f"Error initializing model: {e}")
raise
def create_review_prompt(self, code: str, language: str) -> str:
"""Create a structured prompt for code review."""
return f"""Review this {language} code. List specific points in these sections:
Issues:
Improvements:
Best Practices:
Security:
Code:
```{language}
{code}
```"""
def review_code(self, code: str, language: str) -> str:
"""Perform code review using the model."""
try:
prompt = self.create_review_prompt(code, language)
inputs = self.tokenizer(
prompt,
return_tensors="pt",
truncation=True,
max_length=512,
padding=True
)
with torch.no_grad():
outputs = self.model.generate(
**inputs,
max_new_tokens=512,
do_sample=True,
temperature=0.7,
top_p=0.95,
num_beams=1,
early_stopping=True
)
response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
return response[len(prompt):].strip()
except Exception as e:
logger.error(f"Error during code review: {e}")
return f"Error performing code review: {str(e)}"
# Initialize the reviewer
reviewer = CodeReviewer()
def review_code_interface(code: str, language: str) -> str:
"""Gradio interface function for code review."""
if not code.strip():
return "Please enter some code to review."
try:
result = reviewer.review_code(code, language)
return result
except Exception as e:
return f"Error: {str(e)}"
# Create Gradio interface
iface = gr.Interface(
fn=review_code_interface,
inputs=[
gr.Textbox(
lines=10,
placeholder="Enter your code here...",
label="Code"
),
gr.Dropdown(
choices=["python", "javascript", "java", "cpp", "typescript", "go", "rust"],
value="python",
label="Language"
)
],
outputs=gr.Textbox(
label="Review Results",
lines=10
),
title="Code Review Assistant",
description="An automated code review system powered by Gemma-2b that provides intelligent code analysis and suggestions for improvements.",
examples=[
["""def add_numbers(a, b):
return a + b""", "python"],
["""function calculateSum(numbers) {
let sum = 0;
for(let i = 0; i < numbers.length; i++) {
sum += numbers[i];
}
return sum;
}""", "javascript"]
],
theme=gr.themes.Soft()
)
# Launch the app
if __name__ == "__main__":
iface.launch()
|