Update app_v3.py

app_v3.py

@@ -18,41 +18,24 @@ model = AutoGPTQForCausalLM.from_quantized(model_name_or_path,
         quantize_config=None)
 
 
-prompt = "Tell me about AI"
-prompt_template=f'''### HUMAN:
-{prompt}
-
-### RESPONSE:
-'''
-print("\n\n*** Generate:")
-start_time = time.time()
-input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids.cuda()
-streamer = TextStreamer(tokenizer)
-# output = model.generate(inputs=input_ids, temperature=0.7, max_new_tokens=512)
-# print(tokenizer.decode(output[0]))
-
-_ = model.generate(inputs=input_ids, streamer=streamer, temperature=0.7, max_new_tokens=512)
-print(f"Inference time: {time.time() - start_time:.4f} seconds")
-
-# Inference can also be done using transformers' pipeline
-
-# Prevent printing spurious transformers error when using pipeline with AutoGPTQ
-logging.set_verbosity(logging.CRITICAL)
-
-print("*** Pipeline:")
-start_time = time.time()
-
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    streamer=streamer,
-    max_new_tokens=512,
-    temperature=0.7,
-    top_p=0.95,
-    repetition_penalty=1.15
 )
-
-pipe(prompt_template)
-#print(pipe(prompt_template)[0]['generated_text'])
-print(f"Inference time: {time.time() - start_time:.4f} seconds")
+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'])