Update app.py

app.py

@@ -2,52 +2,45 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Load a lightweight model
-model_name = "mrm8488/GPT-2-finetuned-on-conversational-data"  # ~500MB
+# Use a lightweight and public model
+model_name = "distilgpt2"  # You can also use "tiiuae/falcon-rw-1b" or "EleutherAI/gpt-neo-1.3B"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-# Response generation logic
-def generate_response(message, chat_history):
-    try:
-        chat_history = chat_history or []
-        full_prompt = ""
-        for turn in chat_history:
-            full_prompt += f"User: {turn['content']}\n"
-            full_prompt += f"AI: {turn.get('response', '')}\n"
-        full_prompt += f"User: {message}\nAI:"
-
-        inputs = tokenizer.encode(full_prompt, return_tensors="pt").to(device)
-        outputs = model.generate(inputs, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
-        decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        # Extract only the latest reply after "AI:"
-        response = decoded.split("AI:")[-1].strip()
-
-        # Add to chat history
-        chat_history.append({"role": "user", "content": message})
-        chat_history.append({"role": "assistant", "content": response})
-        return chat_history
-    except Exception as e:
-        return chat_history + [{"role": "assistant", "content": f"⚠️ Error: {str(e)}"}]
-
-# Gradio interface
-with gr.Blocks(css="body { background-color: #FFF9C4; }") as demo:
-    gr.Markdown("<h1 style='color:brown; font-weight:bold;'>🧠 AI Chatbot</h1>")
-    chatbot = gr.Chatbot(label="Talk to AI", type="messages")
-    msg = gr.Textbox(label="Ask something...", placeholder="Type your message here...")
-
-    clear = gr.Button("Clear Chat")
-    footer = gr.Markdown("<p style='text-align: right; color: brown;'>Designed by Mehak Mazhar</p>")
-
-    state = gr.State([])
-
-    msg.submit(generate_response, [msg, state], [chatbot, state])
-    clear.click(lambda: ([], []), None, [chatbot, state])
-
-# Run
+# Define text generation function
+def generate_response(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(
+        inputs["input_ids"],
+        max_length=100,
+        pad_token_id=tokenizer.eos_token_id,
+        do_sample=True,
+        top_k=50,
+        top_p=0.95,
+        temperature=0.7,
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Gradio interface with styling
+def build_interface():
+    with gr.Blocks(theme=gr.themes.Base(), css="""
+        body { background-color: #FFFACD; }
+        h1 { color: brown; font-weight: bold; text-align: center; }
+        footer { text-align: center; padding-top: 10px; font-style: italic; color: #555; }
+    """) as demo:
+        gr.Markdown("# AI Text Generation Chatbot")
+        with gr.Row():
+            with gr.Column():
+                input_text = gr.Textbox(label="Enter your prompt", placeholder="e.g., Once upon a time...")
+                submit_btn = gr.Button("Generate Text")
+            with gr.Column():
+                output_text = gr.Textbox(label="Generated Text")
+
+        submit_btn.click(fn=generate_response, inputs=input_text, outputs=output_text)
+        gr.Markdown("<footer>Designed by Mehak Mazhar</footer>")
+    return demo
+
+# Launch app
 if __name__ == "__main__":
+    demo = build_interface()
     demo.launch()