rephrase / app_v3.py
rodrigomasini's picture
Update app_v3.py
9ebdc85
raw
history blame
1.2 kB
from transformers import AutoTokenizer, TextStreamer, pipeline
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
import time
model_name_or_path = "FPHam/Jackson_The_Formalizer_V2_13b_GPTQ"
model_basename = "Jackson2-gptq_model-4bit-128g"
use_triton = False
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True, legacy=False)
model = AutoGPTQForCausalLM.from_quantized(model_name_or_path,
model_basename=model_basename,
use_safetensors=True,
trust_remote_code=True,
device="cuda:0",
use_triton=use_triton,
quantize_config=None)
)
user_input = st.text_input("Input a phrase")
prompt_template = f'USER: {user_input}\nASSISTANT:'
if st.button("Generate the prompt"):
inputs_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
streamer = TextStreamer(tokenizer)
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
streamer=streamer,
max_new_tokens=512,
temperature=0.2,
top_p=0.95,
repetition_penalty=1.15
)
pipe(prompt_template)
st.write(pipe(prompt_template)[0]['generated_text'])