Update app.py

app.py

@@ -1,60 +1,50 @@
 import gradio as gr
-from transformers import pipeline
-
-# Load text generation models
-models = {
-    "🟦 GPT-2 Tiny": pipeline("text-generation", model="sshleifer/tiny-gpt2"),
-    "🟩 GPT-Neo 125M": pipeline("text-generation", model="EleutherAI/gpt-neo-125M"),
-    "🟨 T5 Small (CommonGen)": pipeline("text2text-generation", model="mrm8488/t5-small-finetuned-common_gen"),
-}
-
-# Generation function
-def generate_text(prompt, model_choice):
-    if not prompt.strip():
-        return "❗ Please enter a prompt."
-    
-    generator = models[model_choice]
-    output = generator(prompt, max_length=100, num_return_sequences=1)[0]
-    
-    return output.get("generated_text", output.get("generated_text", output.get("text", "⚠️ Generation failed.")))
-
-# Gradio App UI
-with gr.Blocks(css="""
-body { background-color: #fffde7; }
-#main-title {
-    text-align: center;
-    font-size: 34px;
-    font-weight: bold;
-    color: #5d4037;
-    margin-top: 20px;
-}
-.output-box {
-    border: 2px solid #ffcc80;
-    padding: 15px;
-    border-radius: 10px;
-    background-color: #fff8e1;
-}
-#footer {
-    text-align: center;
-    font-size: 12px;
-    color: #aaa;
-    margin-top: 30px;
-}
-""") as demo:
-
-    gr.Markdown("<div id='main-title'>📝 LLM for Content Generation</div>")
+from transformers import (
+    GPT2LMHeadModel, GPT2Tokenizer,
+    AutoModelForCausalLM, AutoTokenizer
+)
+
+# Load Models
+gpt2_model = GPT2LMHeadModel.from_pretrained("gpt2")
+gpt2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+
+distilgpt2_model = GPT2LMHeadModel.from_pretrained("distilgpt2")
+distilgpt2_tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
+
+bloom_model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-560m")
+bloom_tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
+
+# Inference Function
+def generate_text(prompt, model_name):
+    if model_name == "🧠 GPT2":
+        inputs = gpt2_tokenizer.encode(prompt, return_tensors="pt")
+        output = gpt2_model.generate(inputs, max_length=100)
+        return gpt2_tokenizer.decode(output[0], skip_special_tokens=True)
+
+    elif model_name == "⚡ DistilGPT2":
+        inputs = distilgpt2_tokenizer.encode(prompt, return_tensors="pt")
+        output = distilgpt2_model.generate(inputs, max_length=100)
+        return distilgpt2_tokenizer.decode(output[0], skip_special_tokens=True)
+
+    elif model_name == "🌸 Bloom-560M":
+        inputs = bloom_tokenizer(prompt, return_tensors="pt")
+        output = bloom_model.generate(inputs["input_ids"], max_length=100)
+        return bloom_tokenizer.decode(output[0], skip_special_tokens=True)
+
+# Gradio UI
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("<h1 style='text-align: center; color: brown;'>LLM for Content Generation</h1>")
+    gr.Markdown("<div style='text-align: center;'>Generate high-quality text using three powerful LLMs</div>")
 
     with gr.Row():
-        with gr.Column(scale=1):
-            prompt_input = gr.Textbox(label="🟨 Prompt", placeholder="Enter your topic, e.g., 'Benefits of AI in education'", lines=5)
-            model_selector = gr.Radio(choices=list(models.keys()), value="🟦 GPT-2 Tiny", label="Select a Model")
-            generate_button = gr.Button("🚀 Generate Content")
+        with gr.Column():
+            prompt = gr.Textbox(label="Enter a topic or prompt")
+            model_choice = gr.Radio(["🧠 GPT2", "⚡ DistilGPT2", "🌸 Bloom-560M"], label="Choose a Model")
+            submit = gr.Button("Generate")
 
-        with gr.Column(scale=1):
-            generated_output = gr.Textbox(label="📝 Generated Text", lines=15, elem_classes=["output-box"])
+        with gr.Column():
+            output = gr.Textbox(label="Generated Text", lines=10)
 
-    generate_button.click(fn=generate_text, inputs=[prompt_input, model_selector], outputs=generated_output)
-
-    gr.Markdown("<div id='footer'>✨ Powered by Hugging Face Transformers</div>")
+    submit.click(fn=generate_text, inputs=[prompt, model_choice], outputs=output)
 
 demo.launch()