Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
+
 class TextGenerationBot:
     def __init__(self, model_name="umairrrkhan/english-text-generation"):
         self.model_name = model_name
@@ -13,16 +14,17 @@ class TextGenerationBot:
     def setup_model(self):
         self.model = AutoModelForCausalLM.from_pretrained(self.model_name)
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        
+
+        # Set pad_token if not defined
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
-            
+
         if self.model.config.pad_token_id is None:
             self.model.config.pad_token_id = self.model.config.eos_token_id
 
     def generate_text(self, input_text, temperature=0.7, max_length=100):
         inputs = self.tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
-        
+
         generation_config = {
             'input_ids': inputs['input_ids'],
             'max_length': max_length,
@@ -35,24 +37,26 @@ class TextGenerationBot:
             'pad_token_id': self.tokenizer.pad_token_id,
             'attention_mask': inputs['attention_mask']
         }
-        
+
         with torch.no_grad():
             outputs = self.model.generate(**generation_config)
-            
+
         return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    def chat(self, message, history):
+    def chat(self, message, history=None):
         self.history = history or []
         bot_response = self.generate_text(message)
         self.history.append((message, bot_response))
         return self.history
 
+
 class ChatbotInterface:
     def __init__(self):
         self.bot = TextGenerationBot()
         self.setup_interface()
 
     def setup_interface(self):
+        # Removed invalid arguments (retry_btn, undo_btn, clear_btn)
         self.interface = gr.ChatInterface(
             fn=self.bot.chat,
             title="AI Text Generation Chatbot",
@@ -62,15 +66,13 @@ class ChatbotInterface:
                 ["What are the benefits of exercise?"],
                 ["Write a poem about nature"]
             ],
-            theme=gr.themes.Soft(),
-            retry_btn="Regenerate Response",
-            undo_btn="Remove Last Message",
-            clear_btn="Clear Conversation",
+            theme=gr.themes.Soft()  # Optional
         )
 
     def launch(self, **kwargs):
         self.interface.launch(**kwargs)
 
+
 def main():
     chatbot = ChatbotInterface()
     chatbot.launch(
@@ -80,5 +82,6 @@ def main():
         debug=True
     )
 
+
 if __name__ == "__main__":
-    main()
+    main()