Update app_v4.py

app_v4.py

@@ -64,7 +64,9 @@ if model_loaded:
     st.write(f"GPU Memory Info after loading the model: {gpu_memory_after}")
 
     # User input for the model
-    user_input = st.text_input("Input a phrase")
+    col1, col2 =st.columns(2)
+    user_input = col1.st.text_input("Input a phrase")
+    max_token = col2.st.number_input(label="Select maxnumber of generated tokens", min_value=1, max_value=1024, value=350, step = 5)
 
     # Generate button
     if st.button("Generate the prompt"):
@@ -74,7 +76,7 @@ if model_loaded:
             inputs = inputs.to(device)  # Move inputs to the same device as model
         # Generate text using torch.inference_mode for better performance during inference
             with torch.inference_mode():
-                output = model.generate(**inputs, max_new_tokens=50)  # Adjust max_new_tokens if needed
+                output = model.generate(**inputs, max_new_tokens=max_token)  
             
             # Cut the tokens at the input length to display only the generated text
             output_ids_cut = output[:, inputs["input_ids"].shape[1]:]