karthik18AI's picture
update with new model
4a0e379 verified
raw
history blame
1.1 kB
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)