change model, log messages

app.py

@@ -1,33 +1,67 @@
 import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import uuid
+import os
+from datetime import datetime
 
-# You can replace this with a file read, environment variable, or UI input
+# ----- Constants -----
+MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
 with open("system_prompt.txt", "r") as f:
-    system_prompt = f.read()
-
-def respond(message, chat_history):
-    if chat_history is None:
-        chat_history = []
-
-    # Combine system prompt and previous messages
-    full_prompt = system_prompt + "\n"
-    for user_msg, bot_msg in chat_history:
-        full_prompt += f"User: {user_msg}\nAssistant: {bot_msg}\n"
-    full_prompt += f"User: {message}\nAssistant:"
-
-    # Use a simple model (or call an API)
-    # Example using Hugging Face transformers (e.g., tiny model for demo):
-    from transformers import pipeline
-    generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.1")
-    response = generator(full_prompt, max_new_tokens=100)[0]['generated_text'].split("Assistant:")[-1].strip()
-
-    chat_history.append((message, response))
-    return "", chat_history
-
-chatbot = gr.Chatbot()
-interface = gr.Interface(fn=respond,
-                         inputs=["text", "state"],
-                         outputs=["text", "state"],
-                         live=False,
-                         title="Custom Prompt Chatbot")
-
-interface.launch()
+    SYSTEM_PROMPT = f.read()
+LOG_DIR = "chat_logs"
+os.makedirs(LOG_DIR, exist_ok=True)
+
+# ----- Load model and tokenizer -----
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto" if device == "cuda" else None
+)
+model.eval()
+
+# ----- Log setup -----
+session_id = str(uuid.uuid4())
+
+def log_chat(session_id, user_msg, bot_msg):
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    with open(os.path.join(LOG_DIR, f"{session_id}.txt"), "a") as f:
+        f.write(f"[{timestamp}] User: {user_msg}\n")
+        f.write(f"[{timestamp}] Bot: {bot_msg}\n\n")
+
+# ----- Inference -----
+def format_chat_prompt(history, new_input):
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    for user_msg, bot_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": bot_msg})
+    messages.append({"role": "user", "content": new_input})
+    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+@torch.no_grad()
+def respond(message, history):
+    prompt = format_chat_prompt(history, message)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    output = model.generate(
+        **inputs,
+        max_new_tokens=512,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.95,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+    # Extract the assistant's final message
+    response = decoded.split(message)[-1].strip().split("\n")[0].strip()
+    log_chat(session_id, message, response)
+    return response
+
+# ----- Gradio Chat Interface -----
+gr.ChatInterface(
+    fn=respond,
+    title="BoundrAI",
+    theme="soft",  # optional aesthetic
+).launch()