Updating the Tokenizer to TF version

app.py

@@ -6,7 +6,7 @@ import os
 from trainer import train
 from tester import test
 import transformers
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import TFAutoModelForCausalLM, TFAutoTokenizer
 
 
 def main():
@@ -34,9 +34,10 @@ def main():
 
 def perform_training(jammer_type, channel_switching_cost):
     agent = train(jammer_type, channel_switching_cost)
+    st.subheader("Generating Insights of the DRL-Training")
     model_name = "tiiuae/falcon-7b-instruct"
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = TFAutoModelForCausalLM.from_pretrained(model_name)
+    tokenizer = TFAutoTokenizer.from_pretrained(model_name)
     pipeline = transformers.pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=100, temperature=0.7)
     text = pipeline("Discuss this topic: Integrating LLMs to DRL-based anti-jamming.")
     st.write(text)

trainer.py

@@ -10,7 +10,7 @@ from antiJamEnv import AntiJamEnv
 
 
 def train(jammer_type, channel_switching_cost):
-    st.subheader("Training Progress")
+    st.subheader("DRL Training Progress")
     progress_bar = st.progress(0)
     status_text = st.empty()
 
@@ -62,7 +62,7 @@ def train(jammer_type, channel_switching_cost):
         if len(rewards) > 10 and np.average(rewards[-10:]) >= max_env_steps - 0.10 * max_env_steps:
             break
 
-    st.sidebar.success("Training completed!")
+    st.sidebar.success("DRL Training completed!")
 
     # Plotting
     rolling_average = np.convolve(rewards, np.ones(10) / 10, mode='valid')