Update app.py

app.py

@@ -1,109 +1,26 @@
 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.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)
-
-        # 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.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)
-
-        # Generate output
-        with torch.no_grad():
-            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):
-        """
-        Handle a chat conversation.
-        """
-        if not history:
-            history = []
-        bot_response = self.generate_text(message)
-        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):
-        """
-        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"],
-            ],
-        )
-
-    def launch(self, **kwargs):
-        """
-        Launch the Gradio interface.
-        """
-        self.interface.launch(**kwargs)
-
-
-def main():
-    chatbot = ChatbotInterface()
-    chatbot.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=True,
-        debug=True,
+# Load the model and tokenizer from Hugging Face
+model_name = "your-username/your-repo-name"  # Replace with your actual model repo name
+model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Define a function to generate text
+def generate_text(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    outputs = model.generate(
+        inputs['input_ids'],
+        max_length=50,
+        attention_mask=inputs['attention_mask'],
+        do_sample=True,
+        temperature=0.7,
+        top_k=50
     )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+# Create a Gradio interface
+iface = gr.Interface(fn=generate_text, inputs="text", outputs="text")
 
-if __name__ == "__main__":
-    main()
+# Launch the interface
+iface.launch()