Update app.py

app.py

@@ -204,7 +204,7 @@ if uploaded_file:
             t_start = time()
             full_prompt = f"<|user|>\nQuestion: {query}\n<|assistant|>"
             input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids
-            st.session_state.input_tokens_count += input_ids.shape[-1]
+            input_tokens_count += input_ids.shape[-1]
             t_end = time()
             log.append(f"✍️ Tokenization Time: {t_end - t_start:.2f} s")
 
@@ -215,8 +215,8 @@ if uploaded_file:
             log.append(f"💡 Generation Time: {last_generation_time:.2f} s")
 
             generated_tokens_count = output_ids.shape[-1]
-            st.session_state.generated_tokens_count += generated_tokens_count
-            st.session_state.output_tokens_count = generated_tokens_count
+            generated_tokens_count += generated_tokens_count
+            output_tokens_count = generated_tokens_count
 
             response = tokenizer.decode(output_ids[0], skip_special_tokens=True)