Replaced review_code method to resolve token_type_ids issue.
Browse files- lib/code_reviewer.py +3 -1
lib/code_reviewer.py
CHANGED
@@ -67,7 +67,9 @@ class CodeReviewer:
|
|
67 |
str: The review generated by the model.
|
68 |
"""
|
69 |
prompt = self.generate_prompt(code)
|
70 |
-
inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
|
|
|
|
|
71 |
output = self.model.generate(**inputs, max_length=512)
|
72 |
review_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
|
73 |
return review_text
|
|
|
67 |
str: The review generated by the model.
|
68 |
"""
|
69 |
prompt = self.generate_prompt(code)
|
70 |
+
inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True).to(self.device)
|
71 |
+
# Remove unsupported keys from model input
|
72 |
+
inputs = {k: v for k, v in inputs.items() if k in self.model.forward.__code__.co_varnames}
|
73 |
output = self.model.generate(**inputs, max_length=512)
|
74 |
review_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
|
75 |
return review_text
|