Spaces:

alKoGolik
/

codellama-CodeLlama-7b-hf

Running

App Files Files Community

codellama-CodeLlama-7b-hf / Llama2-Code-Interpreter-main /code_interpreter /JuypyterClient.py

alKoGolik

Upload 169 files

c87c295 verified 4 months ago

raw history blame contribute delete

No virus

2.51 kB

	from jupyter_client import KernelManager
	import threading
	import re


	class JupyterNotebook:
	def __init__(self):
	self.km = KernelManager()
	self.km.start_kernel()
	self.kc = self.km.client()

	def clean_output(self, outputs):
	outputs_only_str = list()
	for i in outputs:
	if type(i) == dict:
	if "text/plain" in list(i.keys()):
	outputs_only_str.append(i["text/plain"])
	elif type(i) == str:
	outputs_only_str.append(i)
	elif type(i) == list:
	error_msg = "\n".join(i)
	error_msg = re.sub(r"\x1b\[.*?m", "", error_msg)
	outputs_only_str.append(error_msg)

	return "\n".join(outputs_only_str).strip()

	def add_and_run(self, code_string):
	# This inner function will be executed in a separate thread
	def run_code_in_thread():
	nonlocal outputs, error_flag

	# Execute the code and get the execution count
	msg_id = self.kc.execute(code_string)

	while True:
	try:
	msg = self.kc.get_iopub_msg(timeout=20)

	msg_type = msg["header"]["msg_type"]
	content = msg["content"]

	if msg_type == "execute_result":
	outputs.append(content["data"])
	elif msg_type == "stream":
	outputs.append(content["text"])
	elif msg_type == "error":
	error_flag = True
	outputs.append(content["traceback"])

	# If the execution state of the kernel is idle, it means the cell finished executing
	if msg_type == "status" and content["execution_state"] == "idle":
	break
	except:
	break

	outputs = []
	error_flag = False

	# Start the thread to run the code
	thread = threading.Thread(target=run_code_in_thread)
	thread.start()

	# Wait for 10 seconds for the thread to finish
	thread.join(timeout=10)

	# If the thread is still alive after 10 seconds, it's a timeout
	if thread.is_alive():
	outputs = ["Timeout after 10 seconds"]
	error_flag = True

	return self.clean_output(outputs), error_flag

	def close(self):
	"""Shutdown the kernel."""
	self.km.shutdown_kernel()