Update app.py

app.py

@@ -2,15 +2,15 @@ import gradio as gr
 from transformers import pipeline
 import time
 
-# Load the text generation model (GPT-2)
-text_generator = pipeline('text-generation', model='gpt2')
+# Load the smaller and faster text generation model (DistilGPT-2)
+text_generator = pipeline('text-generation', model='distilgpt2')
 
 def get_completion(code_snippet):
     print("Received code snippet, starting processing...")
     start_time = time.time()
-    
+
+    # Provide simplified examples
     python_code_examples = """
-    ---------------------
     Example 1:
     x = 10
     def foo():
@@ -18,68 +18,30 @@ def get_completion(code_snippet):
         x = 5
     foo()
     print(x)
-    Correct output: 5
-    Explanation: The global keyword allows foo() to modify the global variable x, so print(x) outputs 5.
+    Expected output: 5
+    Explanation: `global` allows `foo` to modify `x`.
 
-    ---------------------
     Example 2:
-    def modify_list(input_list):
-        input_list.append(4)
-        input_list = [1, 2, 3]
-    my_list = [0]
-    modify_list(my_list)
-    print(my_list)
-    Correct output: [0, 4]
-    Explanation: The function appends 4 to input_list, which modifies my_list directly. The reassignment does not affect the original list.
-
-    ---------------------
-    Example 3:
     print('hello')
-    Correct output: hello
-    Explanation: This line outputs the string 'hello' to the console.
-    ---------------------
-    Example 4:
-    for i in range(5):
-        print(i)
-    Correct output: 
-    0
-    1
-    2
-    3
-    4
-    Explanation: This loop iterates from 0 to 4, printing each number on a new line.
-    ---------------------
-    Example 5:
-    def add(a, b):
-        return a + b
-    result = add(3, 4)
-    print(result)
-    Correct output: 7
-    Explanation: The function add takes two arguments, adds them together, and returns the result. print(result) outputs 7.
-    ---------------------
-    Example 6:
-    if True:
-        print("This is true")
-    Correct output: This is true
-    Explanation: The condition True always evaluates to True, so the print statement is executed.
-    ---------------------
+    Expected output: hello
+    Explanation: This line prints `hello`.
     """
-    
+
     prompt = f"""
-    Your task is to explain the following code snippet step-by-step.
-    Use the following examples as references for the explanations:
+    Explain the following code snippet step-by-step using these examples as references:
     {python_code_examples}
     Code Snippet:
     ```
     {code_snippet}
     ```
+    Provide a detailed breakdown and expected output.
     """
 
     try:
         print("Generating text...")
-        response = text_generator(prompt, max_new_tokens=100, num_return_sequences=1)
+        response = text_generator(prompt, max_new_tokens=50, num_return_sequences=1)
         print("Text generation complete.")
-        return response[0]['generated_text']
+        return response[0]['generated_text'].strip()
     except Exception as e:
         return f"An error occurred: {str(e)}"
     finally: