|
import sys
|
|
import os
|
|
|
|
prj_root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.append(prj_root_path)
|
|
|
|
from code_interpreter.BaseCodeInterpreter import BaseCodeInterpreter
|
|
from utils.const import *
|
|
|
|
from typing import List, Tuple, Dict
|
|
import re
|
|
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
|
|
sys.path.append(os.path.dirname(__file__))
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
import warnings
|
|
|
|
warnings.filterwarnings("ignore", category=UserWarning, module="transformers")
|
|
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
|
|
|
|
|
|
class OpenCodeInterpreter(BaseCodeInterpreter):
|
|
def __init__(
|
|
self,
|
|
model_path: str,
|
|
load_in_8bit: bool = False,
|
|
load_in_4bit: bool = False,
|
|
):
|
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained(
|
|
model_path,
|
|
padding_side="right",
|
|
trust_remote_code=True
|
|
)
|
|
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
model_path,
|
|
device_map="auto",
|
|
load_in_4bit=load_in_4bit,
|
|
load_in_8bit=load_in_8bit,
|
|
torch_dtype=torch.float16,
|
|
trust_remote_code=True
|
|
)
|
|
|
|
self.model.resize_token_embeddings(len(self.tokenizer))
|
|
|
|
self.model = self.model.eval()
|
|
|
|
self.dialog = []
|
|
self.MAX_CODE_OUTPUT_LENGTH = 1000
|
|
|
|
|
|
def dialog_to_prompt(self, dialog: List[Dict]) -> str:
|
|
full_str = self.tokenizer.apply_chat_template(dialog, tokenize=False)
|
|
|
|
return full_str
|
|
|
|
def extract_code_blocks(self, prompt: str) -> Tuple[bool, str]:
|
|
pattern = re.escape("```python") + r"(.*?)" + re.escape("```")
|
|
matches = re.findall(pattern, prompt, re.DOTALL)
|
|
|
|
if matches:
|
|
|
|
return True, matches[-1].strip()
|
|
else:
|
|
return False, ""
|
|
|
|
def clean_code_output(self, output: str) -> str:
|
|
if self.MAX_CODE_OUTPUT_LENGTH < len(output):
|
|
return (
|
|
output[: self.MAX_CODE_OUTPUT_LENGTH // 5]
|
|
+ "\n...(truncated due to length)...\n"
|
|
+ output[-self.MAX_CODE_OUTPUT_LENGTH // 5 :]
|
|
)
|
|
|
|
return output
|
|
|
