Update app_v1.py

app_v1.py

@@ -6,7 +6,7 @@ import os
 import torch
 
 # Clear up some memory
-# torch.cuda.empty_cache()
+#torch.cuda.empty_cache()
 
 # Try reducing the number of threads PyTorch uses
 # torch.set_num_threads(1)
@@ -54,22 +54,21 @@ model = AutoGPTQForCausalLM.from_quantized(
     quantize_config=quantize_config
 )
 
-st.write(model.hf_device_map)
+#st.write(model.hf_device_map)
+user_input = st.text_input("Input a phrase")
 
-#user_input = st.text_input("Input a phrase")
-
-#prompt_template = f'USER: {user_input}\nASSISTANT:'
+prompt_template = f'USER: {user_input}\nASSISTANT:'
 
 # Generate output when the "Generate" button is pressed
-#if st.button("Generate the prompt"):
-#    inputs = tokenizer(prompt_template, return_tensors="pt")
-#    outputs = model.generate(
-#        input_ids=inputs.input_ids.to("cuda:0"),
-#        attention_mask=inputs.attention_mask.to("cuda:0"),
-#        max_length=512 + inputs.input_ids.size(-1),
-#        temperature=0.1,
-#        top_p=0.95,
-#        repetition_penalty=1.15
-#   )
-#    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-#    st.text_area("Prompt", value=generated_text)
+if st.button("Generate the prompt"):
+    inputs = tokenizer(prompt_template, return_tensors="pt")
+    outputs = model.generate(
+        input_ids=inputs.input_ids.to("cuda:0"),
+        attention_mask=inputs.attention_mask.to("cuda:0"),
+        max_length=512 + inputs.input_ids.size(-1),
+        temperature=0.1,
+        top_p=0.95,
+        repetition_penalty=1.15
+   )
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    st.text_area("Prompt", value=generated_text)