Code_Review_Agent / code_reviewer.py
kumar1907's picture
Update code_reviewer.py
91d1023 verified
raw
history blame contribute delete
900 Bytes
from transformers import pipeline
# Load once
reviewer = pipeline("text2text-generation", model="google/flan-t5-base")
def review_code(code_text):
# Ask model for both feedback + improved code
prompt = f"""
You are a code review assistant. Analyze the following code diff or snippet:
{code_text}
1. Give clear review feedback with improvements.
2. Then provide a refactored/improved version of the code.
Format the response as:
FEEDBACK:
<your feedback here>
CODE:
<improved code here>
"""
output = reviewer(prompt, max_length=800, do_sample=False)
response = output[0]['generated_text']
# Split into feedback & code
if "CODE:" in response:
feedback, code = response.split("CODE:", 1)
feedback = feedback.replace("FEEDBACK:", "").strip()
return feedback, code.strip()
else:
return response, ""