Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-11-26 01:30:20 +01:00
Code Issues Packages Projects Releases Wiki Activity

Remove T5 support (it sucks)

Browse Source
This commit is contained in:
oobabooga 2023-01-10 23:39:50 -03:00
parent 89fd0180b7
commit b2a2ddcb15
1 changed files with 2 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
server.py
View File

@ -8,8 +8,7 @@ from pathlib import Path
import gradio as gr import gradio as gr
import transformers import transformers
from html_generator import * from html_generator import *
from transformers import AutoTokenizer, T5Tokenizer from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import AutoModelForCausalLM, T5ForConditionalGeneration
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
@ -37,8 +36,6 @@ def load_model(model_name):
model = torch.load(Path(f"torch-dumps/{model_name}.pt")) model = torch.load(Path(f"torch-dumps/{model_name}.pt"))
elif model_name.lower().startswith(('gpt-neo', 'opt-', 'galactica')) and any(size in model_name.lower() for size in ('13b', '20b', '30b')): elif model_name.lower().startswith(('gpt-neo', 'opt-', 'galactica')) and any(size in model_name.lower() for size in ('13b', '20b', '30b')):
model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), device_map='auto', load_in_8bit=True) model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), device_map='auto', load_in_8bit=True)
elif model_name in ['flan-t5', 't5-large']:
model = T5ForConditionalGeneration.from_pretrained(Path(f"models/{model_name}")).cuda()
else: else:
model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda() model = AutoModelForCausalLM.from_pretrained(Path(f"models/{model_name}"), low_cpu_mem_usage=True, torch_dtype=torch.float16).cuda()
@ -46,10 +43,6 @@ def load_model(model_name):
else: else:
settings = ["low_cpu_mem_usage=True"] settings = ["low_cpu_mem_usage=True"]
cuda = "" cuda = ""
if model_name in ['flan-t5', 't5-large']:
command = f"T5ForConditionalGeneration.from_pretrained"
else:
command = "AutoModelForCausalLM.from_pretrained" command = "AutoModelForCausalLM.from_pretrained"
if args.cpu: if args.cpu:
@ -72,8 +65,6 @@ def load_model(model_name):
# Loading the tokenizer # Loading the tokenizer
if model_name.lower().startswith('gpt4chan') and Path(f"models/gpt-j-6B/").exists(): if model_name.lower().startswith('gpt4chan') and Path(f"models/gpt-j-6B/").exists():
tokenizer = AutoTokenizer.from_pretrained(Path("models/gpt-j-6B/")) tokenizer = AutoTokenizer.from_pretrained(Path("models/gpt-j-6B/"))
elif model_name in ['flan-t5', 't5-large']:
tokenizer = T5Tokenizer.from_pretrained(Path(f"models/{model_name}/"))
else: else:
tokenizer = AutoTokenizer.from_pretrained(Path(f"models/{model_name}/")) tokenizer = AutoTokenizer.from_pretrained(Path(f"models/{model_name}/"))