Update app.py

app.py

@@ -2,74 +2,96 @@ 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
         self.model = None
         self.tokenizer = None
-        self.history = []
         self.setup_model()
 
     def setup_model(self):
+        """
+        Load the model and tokenizer, and ensure pad_token and pad_token_id are set.
+        """
         self.model = AutoModelForCausalLM.from_pretrained(self.model_name)
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
 
-        # Set pad_token if not defined
+        # Ensure tokenizer has a pad token
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
 
+        # Ensure model config has pad_token_id
         if self.model.config.pad_token_id is None:
-            self.model.config.pad_token_id = self.model.config.eos_token_id
+            self.model.config.pad_token_id = self.tokenizer.pad_token_id
 
     def generate_text(self, input_text, temperature=0.7, max_length=100):
+        """
+        Generate text based on user input.
+        """
+        # Tokenize input
         inputs = self.tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
 
-        generation_config = {
-            'input_ids': inputs['input_ids'],
-            'max_length': max_length,
-            'num_return_sequences': 1,
-            'no_repeat_ngram_size': 2,
-            'temperature': temperature,
-            'top_p': 0.95,
-            'top_k': 50,
-            'do_sample': True,
-            'pad_token_id': self.tokenizer.pad_token_id,
-            'attention_mask': inputs['attention_mask']
-        }
-
+        # Generate output
         with torch.no_grad():
-            outputs = self.model.generate(**generation_config)
-
+            outputs = self.model.generate(
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                max_length=max_length,
+                temperature=temperature,
+                top_k=50,
+                top_p=0.95,
+                do_sample=True,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+            )
+
+        # Decode and return the generated text
         return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    def chat(self, message, history=None):
-        self.history = history or []
+    def chat(self, message, history):
+        """
+        Handle a chat conversation.
+        """
+        if not history:
+            history = []
         bot_response = self.generate_text(message)
-        self.history.append((message, bot_response))
-        return self.history
+        history.append((message, bot_response))
+        return history, history
 
 
 class ChatbotInterface:
     def __init__(self):
         self.bot = TextGenerationBot()
+        self.interface = None
         self.setup_interface()
 
     def setup_interface(self):
-        # Removed invalid arguments (retry_btn, undo_btn, clear_btn)
-        self.interface = gr.ChatInterface(
+        """
+        Set up the Gradio interface for the chatbot.
+        """
+        self.interface = gr.Interface(
             fn=self.bot.chat,
+            inputs=[
+                gr.inputs.Textbox(label="Your Message"),
+                gr.inputs.State(label="Chat History"),
+            ],
+            outputs=[
+                gr.outputs.Textbox(label="Bot Response"),
+                gr.outputs.State(label="Updated Chat History"),
+            ],
             title="AI Text Generation Chatbot",
             description="Chat with an AI model trained on English text. Try asking questions or providing prompts!",
             examples=[
                 ["Tell me a short story about a brave knight"],
                 ["What are the benefits of exercise?"],
-                ["Write a poem about nature"]
+                ["Write a poem about nature"],
             ],
-            theme=gr.themes.Soft()  # Optional
         )
 
     def launch(self, **kwargs):
+        """
+        Launch the Gradio interface.
+        """
         self.interface.launch(**kwargs)
 
 
@@ -79,7 +101,7 @@ def main():
         server_name="0.0.0.0",
         server_port=7860,
         share=True,
-        debug=True
+        debug=True,
     )