CodeGeneration / app.py
Rajut's picture
Update app.py
8d9d71b verified
raw
history blame contribute delete
No virus
2.15 kB
import streamlit as st
from transformers import pipeline
import csv
import re
import torch
import warnings
warnings.filterwarnings("ignore")
# Define a prompt template for Magicoder with placeholders for instruction and response.
MAGICODER_PROMPT = """You are an exceptionally intelligent coding assistant that consistently delivers accurate and reliable responses to user instructions.
@@ Instruction
{instruction}
@@ Response
"""
@st.cache(allow_output_mutation=True)
def load_model():
return pipeline(
model="ise-uiuc/Magicoder-S-DS-6.7B",
task="text-generation"
)
# Function to generate response
def generate_response(instruction):
prompt = MAGICODER_PROMPT.format(instruction=instruction)
result = model(prompt, max_length=2048, num_return_sequences=1, temperature=0.0)
response = result[0]["generated_text"]
response_start_index = response.find("@@ Response") + len("@@ Response")
response = response[response_start_index:].strip()
return response
# Function to append data to a CSV file
def save_to_csv(data, filename):
with open(filename, 'a', newline='') as csvfile:
writer = csv.writer(csvfile)
writer.writerow(data)
# Streamlit app
def main():
global model
st.title("Magicoder Assistant")
if 'model' not in globals():
model = load_model()
instruction = st.text_area("Enter your instruction here:")
if st.button("Generate Response"):
generated_response = generate_response(instruction)
st.text("Generated response:")
st.text(generated_response)
correct_output = st.radio("Is the generated output correct?", ("Yes", "No"))
if correct_output.lower() == 'yes':
feedback = st.text_input("Do you want to provide any feedback?")
save_to_csv(["Correct", feedback], 'output_ratings.csv')
else:
correct_code = st.text_area("Please enter the correct code:")
feedback = st.text_input("Any other feedback you want to provide:")
save_to_csv(["Incorrect", feedback, correct_code], 'output_ratings.csv')
if __name__ == "__main__":
main()