Update app.py

app.py

@@ -1,8 +1,8 @@
 import streamlit as st
 from transformers import pipeline
-import torch
 import csv
 import re
+import torch
 import warnings
 
 warnings.filterwarnings("ignore")
@@ -14,18 +14,17 @@ MAGICODER_PROMPT = """You are an exceptionally intelligent coding assistant that
 @@ Response
 """
 
-# Create a text generation pipeline using the Magicoder model, text-generation task, bfloat16 torch data type and auto device mapping.
-generator = pipeline(
-    model="ise-uiuc/Magicoder-S-DS-6.7B",
-    task="text-generation",
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-)
+@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 = generator(prompt, max_length=2048, num_return_sequences=1, temperature=0.0)
+    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()
@@ -37,20 +36,14 @@ def save_to_csv(data, filename):
         writer = csv.writer(csvfile)
         writer.writerow(data)
 
-# Function to process user feedback
-def process_output(correct_output):
-    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')
-
 # 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)
@@ -58,7 +51,13 @@ def main():
         st.text(generated_response)
 
         correct_output = st.radio("Is the generated output correct?", ("Yes", "No"))
-        process_output(correct_output)
+        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()