Update app.py

app.py

@@ -1,26 +1,63 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import logging
 
-# Load the model and tokenizer from Hugging Face
-model_name = "umairrrkhan/english-text-generation"  # Replace with your actual model repo name
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Define a function to generate text
+# Model and tokenizer setup
+def setup_model_and_tokenizer():
+    logger.info("Loading model and tokenizer...")
+    model_name = "umairrrkhan/english-text-generation"  # Replace with your model
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    # Ensure pad_token is set
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    if model.config.pad_token_id is None:
+        model.config.pad_token_id = tokenizer.pad_token_id
+
+    logger.info("Model and tokenizer loaded successfully.")
+    return model, tokenizer
+
+model, tokenizer = setup_model_and_tokenizer()
+
+# Define text generation function
 def generate_text(prompt):
+    logger.info(f"Received prompt: {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")
+
+    try:
+        logger.info("Generating text...")
+        outputs = model.generate(
+            inputs['input_ids'],
+            max_length=50,
+            attention_mask=inputs['attention_mask'],
+            do_sample=True,
+            temperature=0.7,
+            top_k=50,
+        )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        logger.info(f"Generated response: {response}")
+        return response
+    except Exception as e:
+        logger.error(f"Error during text generation: {e}")
+        return "An error occurred during text generation."
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_text,
+    inputs="text",
+    outputs="text",
+    title="AI Text Generation Chatbot",
+    description="Type a prompt and see what the AI generates!",
+    examples=["Tell me a story about a robot.", "Write a poem about the moon."]
+)
 
 # Launch the interface
-iface.launch()
+if __name__ == "__main__":
+    logger.info("Launching Gradio interface...")
+    iface.launch(debug=True, server_name="0.0.0.0", server_port=7860)