e

app.py

@@ -30,6 +30,7 @@ logging.info(C("[INFO] "f"device = {device}"))
 )
 def model_init():
     
+    logging.info(C("[INFO] "f"Model init start!"))
     
     
     from transformers import GenerationConfig
@@ -65,15 +66,15 @@ try:
 
     # =================== INFERENCE ==================== #
     if prompt:
-        st.balloons()
+        st.title(prompt)
         with torch.no_grad():
             [texts_out] = model.generate(
                 **tokenizer(
                     prompt, return_tensors="pt",
                     
                 ).to(device),
-                max_new_tokens=200,
-                )
+                max_new_tokens=100,
+            )
         output_text = tokenizer.decode(texts_out)
         st.balloons()
         st.markdown(output_text)