vsagar100 commited on
Commit
a9a2195
·
verified ·
1 Parent(s): 0dbfe24

Replaced review_code method to resolve token_type_ids issue.

Browse files
Files changed (1) hide show
  1. 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