【LLM Compiler】Meta社がコード最適化&コンパイル用LLMを開発

# 必要なライブラリのインストール
!pip install transformers
!pip install torch
!pip install sentencepiece
!pip install huggingface_hub
!pip install accelerate
!pip install transformers accelerate

from transformers import AutoTokenizer
import transformers
import torch
import os
from huggingface_hub import login

# Hugging Face APIトークンの設定
token = "your_token"
login(token=token)
os.environ["HUGGING_FACE_HUB_TOKEN"] = token
model = "facebook/llm-compiler-7b"

tokenizer = AutoTokenizer.from_pretrained(model, token=token)
pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    torch_dtype=torch.float16,
    device_map="auto",
    token=token
)

sequences = pipeline(
    '%3 = alloca i32, align 4',
    do_sample=True,
    top_k=10,
    temperature=0.1,
    top_p=0.95,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
    max_length=200,
)

for seq in sequences:
    print(f"Result: {seq['generated_text']}")

Result: %3 = alloca i32, align 4
  %4 = alloca i32, align 4
  %5 = alloca i32, align 4
  %6 = alloca i32, align 4
  %7 = alloca i32, align 4
  %8 = alloca i32, align 4
  %9 = alloca i32, align 4
  %10 = alloca i32, align 4
  %11 = alloca i32, align 4
  %12 = alloca i32, align 4
  %13 = alloca i32, align 4
  %14 = alloca i32, align 4
  %15 = alloca i32, align 4
  %16 = alloca i32, align 

from transformers import AutoTokenizer, pipeline
import torch
from huggingface_hub import login
import os
import re
import time
import random

# Hugging Face APIトークンの設定
token = "Your_token"
login(token=token)
os.environ["HUGGING_FACE_HUB_TOKEN"] = token

# モデルの指定
model = "facebook/llm-compiler-7b"

# トークナイザーとパイプラインの設定
tokenizer = AutoTokenizer.from_pretrained(model, token=token)
pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    torch_dtype=torch.float16,
    device_map="auto",
)

# 最適化前コード
inefficient_code = """
def very_inefficient_sort(arr):
    sorted_arr = []
    while len(arr) > 0:
        min_val = min(arr)
        sorted_arr.append(min_val)
        arr.remove(min_val)
    return sorted_arr
"""

# より明確な最適化の指示を含むプロンプト
input_text = f"""Optimize the following inefficient Python function for sorting an array:

{inefficient_code}

Please provide a more efficient implementation of this sorting function. Your optimized function should:

1. Use either Python's built-in sorted() function or list.sort() method for efficient sorting.
2. Avoid creating unnecessary new lists if possible.
3. Aim for a time complexity of O(n log n) instead of the current O(n^2).

Provide your optimized implementation below:
"""

sequences = pipeline(
    input_text,
    do_sample=True,
    top_k=10,
    temperature=0.7,
    top_p=0.95,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
    max_length=1000,
    truncation=True
)

# 生成されたテキストから関数を抽出
generated_text = sequences[0]['generated_text']
print(f"Generated Text:\n{generated_text}")

# 関数定義を抽出
function_match = re.search(r'def\s+sort_by_mergesort\s*\(.*?\):\s*([\s\S]*?)(?=\ndef|$)', generated_text)

if function_match:
    optimized_code = function_match.group(0)
    print(f"Optimized Code:\n{optimized_code}")
    
    # 最適化されたコードを評価
    try:
        exec(optimized_code, globals())
        print("Optimized function successfully defined.")
        
        # テストデータ
        test_data = [random.randint(1, 1000) for _ in range(1000)]
        
        # 元の関数と最適化された関数を実行して比較
        exec(inefficient_code, globals())
        
        start_time = time.time()
        inefficient_result = very_inefficient_sort(test_data.copy())
        inefficient_time = time.time() - start_time
        print(f"Inefficient Sort Time: {inefficient_time:.6f} seconds")
        
        start_time = time.time()
        optimized_result = sort_by_mergesort(test_data.copy())
        optimized_time = time.time() - start_time
        print(f"Optimized Sort Time: {optimized_time:.6f} seconds")
        
        # 結果の正確性を確認
        print(f"Results match: {inefficient_result == optimized_result}")
        
        # 速度向上の計算
        speedup = inefficient_time / optimized_time
        print(f"Speedup: {speedup:.2f}x")
        
        # より大きなテストデータでの実行
        large_test_data = [random.randint(1, 100000) for _ in range(10000)]
        
        start_time = time.time()
        very_inefficient_sort(large_test_data.copy())
        large_inefficient_time = time.time() - start_time
        print(f"Inefficient Sort Time (large data): {large_inefficient_time:.6f} seconds")
        
        start_time = time.time()
        sort_by_mergesort(large_test_data.copy())
        large_optimized_time = time.time() - start_time
        print(f"Optimized Sort Time (large data): {large_optimized_time:.6f} seconds")
        
        large_speedup = large_inefficient_time / large_optimized_time
        print(f"Speedup (large data): {large_speedup:.2f}x")
        
    except Exception as e:
        print(f"Error in optimized code: {e}")
else:
    print("No valid function found in the generated text. Using the original inefficient code.")

# 最適化前コード
def very_inefficient_sort(arr):
    sorted_arr = []
    while len(arr) > 0:
        min_val = min(arr)
        sorted_arr.append(min_val)
        arr.remove(min_val)
    return sorted_arr

def better_sort(arr):
    arr.sort()
    return arr

import time
import random

# 元の非効率なコード
def very_inefficient_sort(arr):
    sorted_arr = []
    while len(arr) > 0:
        min_val = min(arr)
        sorted_arr.append(min_val)
        arr.remove(min_val)
    return sorted_arr

# 最適化されたコード
def better_sort(arr):
    arr.sort()
    return arr

# テストデータの生成
test_data = [random.randint(1, 1000) for _ in range(1000)]

# 非効率なソートの実行時間を計測
start_time = time.time()
very_inefficient_sort(test_data.copy())
inefficient_time = time.time() - start_time
print(f"Very Inefficient Sort Time: {inefficient_time:.6f} seconds")

# 最適化されたソートの実行時間を計測
start_time = time.time()
better_sort(test_data.copy())
optimized_time = time.time() - start_time
print(f"Better Sort Time: {optimized_time:.6f} seconds")

# 速度向上の計算
speedup = inefficient_time / optimized_time
print(f"Speedup: {speedup:.2f}x")

Very Inefficient Sort Time: 0.042483 seconds
Better Sort Time: 0.000247 seconds
Speedup: 171.66x