Create app.py

app.py

@@ -0,0 +1,57 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from datasets import load_dataset
+import random
+
+# Load tokenizer & model
+tokenizer = AutoTokenizer.from_pretrained("driaforall/mem-agent")
+model = AutoModelForCausalLM.from_pretrained("driaforall/mem-agent")
+
+# Load dataset of prompts
+ds = load_dataset("fka/awesome-chatgpt-prompts")
+prompts = ds["train"]["prompt"]
+
+# Pick a random prompt to initialize Firebox's "persona"
+firebox_persona = random.choice(prompts)
+print("🔥 Firebox Persona:", firebox_persona)
+
+def firebox_chat(user_message, history=[]):
+    """
+    Firebox chatbot using driaforall/mem-agent with persona prompts.
+    """
+    # If history is empty, insert persona/system prompt first
+    if not history:
+        history.append({"role": "system", "content": firebox_persona})
+    
+    # Add user message
+    history.append({"role": "user", "content": user_message})
+
+    # Apply Hugging Face chat template
+    inputs = tokenizer.apply_chat_template(
+        history,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+    ).to(model.device)
+
+    # Generate Firebox reply
+    outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True, temperature=0.7)
+    reply = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+
+    # Add assistant reply to history
+    history.append({"role": "assistant", "content": reply})
+
+    return reply, history
+
+
+# Example usage
+if __name__ == "__main__":
+    history = []
+    print("🔥 Firebox:", firebox_persona)
+    while True:
+        user_input = input("You: ")
+        if user_input.lower() in ["quit", "exit", "bye"]:
+            print("🔥 Firebox: Goodbye! Stay safe.")
+            break
+        reply, history = firebox_chat(user_input, history)
+        print("🔥 Firebox:", reply)