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()