File size: 2,283 Bytes
c87c295 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
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,
):
# build tokenizer
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 the last matched code block
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
|