Update app.py

app.py

@@ -1,34 +1,47 @@
 # app.py
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+
+# -----------------------------
+# 1️⃣ Hugging Face token
+# -----------------------------
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if HF_TOKEN is None:
+    raise ValueError("Set your Hugging Face token in the environment variable HF_TOKEN")
 
-# Hugging Face model repo
 MODEL_REPO = "DSDUDEd/firebase"
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
-model = AutoModelForCausalLM.from_pretrained(MODEL_REPO)
+# -----------------------------
+# 2️⃣ Load model & tokenizer
+# -----------------------------
+tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, use_auth_token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained(MODEL_REPO, use_auth_token=HF_TOKEN)
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-def generate_response(prompt, max_tokens=100):
-    """Generate text from the model."""
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    outputs = model.generate(inputs["input_ids"], max_new_tokens=max_tokens)
+# -----------------------------
+# 3️⃣ Define generation function
+# -----------------------------
+def generate(prompt, max_tokens=50):
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    outputs = model.generate(input_ids, max_new_tokens=max_tokens)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Gradio interface
+# -----------------------------
+# 4️⃣ Create Gradio interface
+# -----------------------------
 iface = gr.Interface(
-    fn=generate_response,
-    inputs=[
-        gr.Textbox(label="Input Prompt"),
-        gr.Slider(minimum=10, maximum=500, step=10, label="Max Tokens")
-    ],
-    outputs=gr.Textbox(label="Model Output"),
-    title="Custom GPT-2 AI",
-    description="Type a prompt and the AI will generate a response."
+    fn=generate,
+    inputs=[gr.Textbox(lines=2, placeholder="Enter your prompt here"), gr.Slider(minimum=1, maximum=200, value=50)],
+    outputs="text",
+    title="DSDUDEd Firebase AI",
+    description="Generate text using the custom model hosted privately on Hugging Face."
 )
 
-if __name__ == "__main__":
-    iface.launch()
+# -----------------------------
+# 5️⃣ Launch Space
+# -----------------------------
+iface.launch()