Update main.py

main.py

@@ -89,6 +89,18 @@ def main():
 
             # Second column: Display generated code
             with col2:
+                def extract_code_from_output(result):
+                if result is not None:
+                output = result.stdout
+                if output is not None and "code>>>:" in output:
+                code = output.split("code>>>:")[1]
+                return code
+                return None  # or handle the case when code extraction is not successful
+
+# Example usage:
+                result = ...  # your result variable
+                code_result = extract_code_from_output(result)
+
                 code_result = extract_code_from_output(result)
                 if option == "gorilla-7b-hf-v1":
                     st.subheader("Output Generated")