Spaces:
Sleeping
Sleeping
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer | |
model_name = "Salesforce/codet5-base" | |
model = AutoModelForSeq2SeqLM.from_pretrained(model_name) | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
def generate_review(code_snippet): | |
inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True, padding="max_length", max_length=512) | |
review_ids = model.generate(inputs.input_ids, max_new_tokens=100, num_beams=4, length_penalty=2.0, early_stopping=True) | |
review = tokenizer.decode(review_ids[0], skip_special_tokens=True) | |
return review | |
# Example C# code snippet | |
code_snippet = """ | |
public class Calculator { | |
public int Add(int a, int b) { | |
return a + b; | |
} | |
public int Subtract(int a, int b) { | |
return a - b; | |
} | |
public int Multiply(int a, int b) { | |
return a * b; | |
} | |
public int Divide(int a, int b) { | |
if (b == 0) { | |
throw new DivideByZeroException("Division by zero is not allowed."); | |
} | |
return a / b; | |
} | |
} | |
""" | |
review = generate_review(code_snippet) | |
print("Code Review:", review) | |