Spaces:
Sleeping
Sleeping
File size: 953 Bytes
33215f4 48e9952 33215f4 48e9952 33215f4 2da07e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
from transformers import pipeline
import streamlit as st
# Load once, cache in Streamlit
@st.cache_resource
def load_model():
return pipeline("text2text-generation", model="google/flan-t5-base")
reviewer = load_model()
def review_code(code_text):
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, ""
|