Update app_v1.py

app_v1.py

@@ -6,10 +6,10 @@ 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)
+# torch.set_num_threads(1)
 
 cwd = os.getcwd()
 cachedir = cwd + '/cache'
@@ -53,20 +53,22 @@ model = AutoGPTQForCausalLM.from_quantized(
     quantize_config=quantize_config
 )
 
-user_input = st.text_input("Input a phrase")
+st.write(model.hf_device_map)
 
-prompt_template = f'USER: {user_input}\nASSISTANT:'
+#user_input = st.text_input("Input a phrase")
+
+#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)