Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import json | |
| import os | |
| import time | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| # === CONFIGURATION === | |
| MODEL_NAME = "facebook/opt-350m" # ✅ Fast and lightweight | |
| MEMORY_FILE = "memory.json" | |
| # === LOAD MODEL === | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_NAME, | |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, | |
| device_map="auto" | |
| ) | |
| # === MEMORY === | |
| def load_memory(file_path): | |
| if os.path.exists(file_path): | |
| with open(file_path, "r", encoding="utf-8") as f: | |
| return json.load(f) | |
| return [] | |
| def save_memory(file_path, data): | |
| with open(file_path, "w", encoding="utf-8") as f: | |
| json.dump(data, f, indent=4) | |
| memory = load_memory(MEMORY_FILE) | |
| # === CHAT FUNCTION === | |
| def chat_with_ai(user_input, history, turbo_mode): | |
| memory.append({"user": user_input}) | |
| save_memory(MEMORY_FILE, memory) | |
| conversation = "" | |
| for msg in memory[-3:]: # Only the last 3 inputs to keep it lightweight | |
| conversation += f"User: {msg['user']}\nAssistant:" | |
| inputs = tokenizer(conversation, return_tensors="pt", truncation=True, max_length=512).to(model.device) | |
| generate_args = { | |
| 'input_ids': inputs.input_ids, | |
| 'max_length': inputs.input_ids.shape[-1] + (30 if turbo_mode else 50), | |
| 'temperature': 0.6, | |
| 'top_k': 30, | |
| 'top_p': 0.9, | |
| 'do_sample': True, | |
| 'pad_token_id': tokenizer.eos_token_id, | |
| 'early_stopping': True | |
| } | |
| output = model.generate(**generate_args) | |
| reply_tokens = output[0][inputs.input_ids.shape[-1]:] | |
| reply_text = tokenizer.decode(reply_tokens, skip_special_tokens=True).strip() | |
| # Save Assistant reply to memory | |
| memory.append({"assistant": reply_text}) | |
| save_memory(MEMORY_FILE, memory) | |
| return "", history + [{"role": "user", "content": user_input}, {"role": "assistant", "content": reply_text}], reply_text | |
| def clear_chat(): | |
| memory.clear() | |
| save_memory(MEMORY_FILE, []) | |
| return "", [], "" | |
| def save_chat(history): | |
| timestamp = time.strftime("%Y-%m-%d_%H-%M-%S") | |
| filename = f"chat_{timestamp}.txt" | |
| with open(filename, "w", encoding="utf-8") as f: | |
| for msg in history: | |
| f.write(f"{msg['role'].capitalize()}: {msg['content']}\n\n") | |
| return f"💾 Chat saved as {filename}." | |
| # === DARK THEME CSS === | |
| css = """ | |
| body { | |
| background-color: #121212; | |
| color: #e0e0e0; | |
| } | |
| .gradio-container, .gr-block, .gr-button, .gr-textbox { | |
| background-color: #1e1e1e !important; | |
| border-color: #333 !important; | |
| } | |
| .gr-button:hover { | |
| background-color: #333333 !important; | |
| } | |
| """ | |
| # === GRADIO APP === | |
| with gr.Blocks(css=css) as app: | |
| gr.Markdown("## 🤖 Ultra Fast Living AI - Clean Mode", elem_id="title") | |
| chatbot = gr.Chatbot(label="Chat History", type="messages", render=False) | |
| user_input = gr.Textbox(lines=3, label="Your Message", placeholder="Type your message...") | |
| ai_output = gr.Textbox(lines=5, label="🤖 AI's Response", interactive=False, placeholder="AI will respond here...") | |
| with gr.Row(): | |
| send_btn = gr.Button("Send") | |
| clear_btn = gr.Button("Clear Chat") | |
| save_btn = gr.Button("Save Chat") | |
| turbo_mode = gr.Checkbox(label="🚀 Turbo Mode (Faster Replies)", value=True) | |
| status_box = gr.Textbox(label="System Status", interactive=False) | |
| send_btn.click(chat_with_ai, inputs=[user_input, chatbot, turbo_mode], outputs=[user_input, chatbot, ai_output]) | |
| user_input.submit(chat_with_ai, inputs=[user_input, chatbot, turbo_mode], outputs=[user_input, chatbot, ai_output]) | |
| clear_btn.click(clear_chat, outputs=[user_input, chatbot, ai_output]) | |
| save_btn.click(save_chat, inputs=[chatbot], outputs=[status_box]) | |
| app.launch() |