Update app_v3.py

app_v3.py

@@ -38,15 +38,15 @@ if torch.cuda.is_available():
     
 #tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True, legacy=False)
 tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 model = AutoGPTQForCausalLM.from_quantized(
     pretrained_model_dir,
     model_basename=model_basename,
     use_safetensors=True,
-    device="cuda:0",
-    #use_triton=use_triton,
-    #quantize_config=None
+    device=device,
 )
 
+
 viz = torch.cuda.memory_summary()
 st.write(viz)
 
@@ -55,7 +55,8 @@ user_input = st.text_input("Input a phrase")
 prompt_template = f'USER: {user_input}\nASSISTANT:'
 
 if st.button("Generate the prompt"):
-    inputs = tokenizer(prompt_template, return_tensors='pt')
+    inputs = tokenizer(prompt_template, return_tensors='pt', max_length=512, truncation=True, padding='max_length', batch_size=4)
+    #inputs = tokenizer(prompt_template, return_tensors='pt')
     #streamer = TextStreamer(tokenizer)
     #pipe = pipeline(
     #    "text-generation",