224 lines
7.2 KiB
Python
224 lines
7.2 KiB
Python
import dotenv
|
|
import os
|
|
from openai import OpenAI
|
|
from anthropic import Anthropic
|
|
import gradio as gr
|
|
|
|
# from .config import *
|
|
|
|
OPENAI_MODEL = "gpt-4o-mini"
|
|
CLAUDE_MODEL = "claude-3-5-haiku-20241022"
|
|
|
|
OUTPUT_MAX_TOKEN = 2000
|
|
|
|
CSS = """
|
|
body {
|
|
background: #f4f6fa;
|
|
font-family: 'Segoe UI', Roboto, sans-serif;
|
|
}
|
|
|
|
.raw textarea {
|
|
border: 1.5px solid #00FFBF !important;
|
|
box-shadow: 0 0 10px rgba(229, 115, 115, 0.3);
|
|
color: #00FFBF !important;
|
|
font-size: 24px;
|
|
}
|
|
|
|
.optimize textarea {
|
|
border: 1.5px solid #FFBF00 !important;
|
|
box-shadow: 0 0 10px rgba(129, 199, 132, 0.3);
|
|
color: #FFBF00 !important;
|
|
font-size: 24px
|
|
}
|
|
|
|
button {
|
|
background: linear-gradient(90deg, #2196f3, #00BFFF);
|
|
color: white !important;
|
|
font-weight: bold;
|
|
border-radius: 8px !important;
|
|
transition: all 0.2s ease-in-out;
|
|
}
|
|
|
|
button:hover {
|
|
background: linear-gradient(90deg, #21cbf3, #2196f3);
|
|
transform: scale(1.05);
|
|
}
|
|
|
|
h1 {
|
|
text-align: center;
|
|
color: #1565c0;
|
|
font-size: 38px;
|
|
}
|
|
"""
|
|
|
|
PYTHON_CODE = '''
|
|
import math
|
|
|
|
def pairwise_distance(points_a, points_b):
|
|
"""
|
|
Compute the pairwise Euclidean distance between two sets of 3D points.
|
|
|
|
Args:
|
|
points_a: list of (x, y, z)
|
|
points_b: list of (x, y, z)
|
|
Returns:
|
|
A 2D list of shape (len(points_a), len(points_b)) representing distances
|
|
"""
|
|
distances = []
|
|
for i in range(len(points_a)):
|
|
row = []
|
|
for j in range(len(points_b)):
|
|
dx = points_a[i][0] - points_b[j][0]
|
|
dy = points_a[i][1] - points_b[j][1]
|
|
dz = points_a[i][2] - points_b[j][2]
|
|
d = math.sqrt(dx * dx + dy * dy + dz * dz)
|
|
row.append(d)
|
|
distances.append(row)
|
|
return distances
|
|
|
|
|
|
# Example usage
|
|
if __name__ == "__main__":
|
|
import random
|
|
points_a = [(random.random(), random.random(), random.random()) for _ in range(100)]
|
|
points_b = [(random.random(), random.random(), random.random()) for _ in range(100)]
|
|
dists = pairwise_distance(points_a, points_b)
|
|
print(f"Distance[0][0] = {dists[0][0]:.4f}")
|
|
'''
|
|
|
|
|
|
def main():
|
|
dotenv.load_dotenv(override=True)
|
|
os.environ['OPENAI_API_KEY'] = os.getenv(
|
|
'OPENAI_API_KEY', 'your-key-if-not-using-env')
|
|
os.environ['ANTHROPIC_API_KEY'] = os.getenv(
|
|
'ANTHROPIC_API_KEY', 'your-key-if-not-using-env')
|
|
|
|
# codeReviser = CodeAccelerator('openai', os.getenv('OPENAI_API_KEY'))
|
|
codeReviser = CodeAccelerator('anthropic', os.getenv('ANTHROPIC_API_KEY'))
|
|
|
|
display_ui(codeReviser)
|
|
|
|
|
|
def safe_exec(code_str):
|
|
import io
|
|
import sys
|
|
import time
|
|
import ast
|
|
# Build the buffer of IO to extract ouput of stdout
|
|
stdout_buffer = io.StringIO()
|
|
old_stdout = sys.stdout
|
|
sys.stdout = stdout_buffer
|
|
|
|
try:
|
|
tree = ast.parse(code_str)
|
|
compiled = compile(tree, filename="<ast>", mode="exec")
|
|
local_vars = {}
|
|
start = time.time()
|
|
exec(compiled, {}, local_vars)
|
|
exec_time = time.time() - start
|
|
print(f"This code spend {exec_time:.8f} seconds\n")
|
|
|
|
# recover sys.stdout
|
|
sys.stdout = old_stdout
|
|
output_text = stdout_buffer.getvalue()
|
|
return output_text
|
|
|
|
except Exception as e:
|
|
sys.stdout = old_stdout
|
|
return f"Error: {e}"
|
|
|
|
|
|
def display_ui(codeReviser):
|
|
def _optimize(pythonCode):
|
|
for text in codeReviser.respond(pythonCode):
|
|
yield text.replace("```python", "").replace("```", "")
|
|
|
|
with gr.Blocks(css=CSS) as ui:
|
|
gr.Markdown("# ✨Convert Python code for accelation")
|
|
with gr.Row():
|
|
beforeBlock = gr.Textbox(
|
|
label="raw python code", value=PYTHON_CODE, lines=20, elem_classes=["raw"])
|
|
afterBlock = gr.Textbox(
|
|
label="optimized python code", lines=20, elem_classes=["optimize"])
|
|
with gr.Row():
|
|
convert = gr.Button("Convert code")
|
|
with gr.Row():
|
|
rawRunButton = gr.Button("Run raw code")
|
|
optRunButton = gr.Button("Run optimized code")
|
|
with gr.Row():
|
|
rawOut = gr.TextArea(label="Raw result:",
|
|
elem_classes=["raw"])
|
|
optimizeOut = gr.TextArea(
|
|
label="Optimize result:", elem_classes=["optimize"])
|
|
|
|
convert.click(_optimize,
|
|
inputs=[beforeBlock], outputs=[afterBlock])
|
|
rawRunButton.click(safe_exec, inputs=[beforeBlock], outputs=[rawOut])
|
|
optRunButton.click(safe_exec, inputs=[
|
|
afterBlock], outputs=[optimizeOut])
|
|
|
|
ui.launch(inbrowser=True)
|
|
|
|
|
|
class CodeAccelerator:
|
|
|
|
def __init__(self, frontierModel: str, apiKey):
|
|
self.frontierModel = frontierModel
|
|
|
|
if frontierModel == "openai":
|
|
self.llm = OpenAI(api_key=apiKey)
|
|
elif frontierModel == "anthropic":
|
|
self.llm = Anthropic(api_key=apiKey)
|
|
else:
|
|
raise ValueError(f'frontierModel {frontierModel} is invalid.')
|
|
|
|
def _getChatTemplate(self, pythonCode):
|
|
_code = pythonCode.strip()
|
|
|
|
systemPrompt = '''
|
|
You are an assistant that reimplements Python code in high performance and just spend the fastest possible time for an windows laptop.
|
|
Respond only with Python code; use comments sparingly and do not provide any explanation other than occasional comments.
|
|
The new Python code response needs to produce an identical output in the fastest possible time.
|
|
'''
|
|
userPrompt = f'''
|
|
Rewrite this Python code with the fastest possible implementation that produces identical output in the least time.
|
|
Respond only with Python code; do not explain your work other than a few comments.
|
|
Remember to import all necessary python packages such as numpy.\n\n
|
|
|
|
{_code}
|
|
'''
|
|
return [
|
|
{"role": "system", "content": systemPrompt},
|
|
{"role": "user", "content": userPrompt}
|
|
]
|
|
|
|
def respond(self, pythonCode):
|
|
"""Generator"""
|
|
chatTemplate = self._getChatTemplate(pythonCode)
|
|
reply = ""
|
|
if self.frontierModel == 'openai':
|
|
stream = self.llm.chat.completions.create(messages=chatTemplate,
|
|
model=OPENAI_MODEL,
|
|
max_tokens=OUTPUT_MAX_TOKEN,
|
|
stream=True)
|
|
for chunk in stream:
|
|
chunkText = chunk.choices[0].delta.content or ""
|
|
reply += chunkText
|
|
yield reply
|
|
elif self.frontierModel == "anthropic":
|
|
stream = self.llm.messages.create(model=CLAUDE_MODEL,
|
|
system=chatTemplate[0]['content'],
|
|
messages=chatTemplate[1:],
|
|
max_tokens=OUTPUT_MAX_TOKEN,
|
|
stream=True)
|
|
|
|
for chunk in stream:
|
|
chunkText = chunk.delta.text if chunk.type == "content_block_delta" else ""
|
|
reply += chunkText
|
|
yield reply
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|