|
|
import os |
|
|
import gradio as gr |
|
|
from transformers import pipeline, Conversation |
|
|
from dotenv import load_dotenv |
|
|
|
|
|
load_dotenv() |
|
|
PASSWORD = os.getenv("ARMAAN_PASS") |
|
|
|
|
|
|
|
|
chatbot = pipeline("conversational", model="facebook/blenderbot-1B-distill") |
|
|
system_prompt = "You are a helpful chatbot." |
|
|
|
|
|
def chat_fn(message, history): |
|
|
convo = Conversation(f"{message}") |
|
|
result = chatbot(convo) |
|
|
return result.generated_responses[-1] |
|
|
|
|
|
def check_password(pass_input): |
|
|
if pass_input == PASSWORD: |
|
|
return gr.update(visible=True), "" |
|
|
else: |
|
|
return gr.update(visible=False), "Incorrect password" |
|
|
|
|
|
def update_prompt(new_prompt): |
|
|
global system_prompt |
|
|
system_prompt = new_prompt |
|
|
return "Prompt updated." |
|
|
|
|
|
with gr.Blocks() as demo: |
|
|
gr.Markdown("# Chatbot (Fast CPU-Friendly)") |
|
|
|
|
|
with gr.Tab("Chat"): |
|
|
gr.ChatInterface(fn=chat_fn) |
|
|
|
|
|
with gr.Tab("Train"): |
|
|
with gr.Column(): |
|
|
pass_input = gr.Textbox(label="Enter Admin Password", type="password") |
|
|
login_btn = gr.Button("Login") |
|
|
error_text = gr.Markdown("", visible=False) |
|
|
|
|
|
with gr.Group(visible=False) as admin_panel: |
|
|
new_prompt = gr.Textbox(label="New System Prompt", lines=4) |
|
|
update_btn = gr.Button("Update Prompt") |
|
|
success_msg = gr.Markdown("") |
|
|
|
|
|
login_btn.click(fn=check_password, inputs=pass_input, outputs=[admin_panel, error_text]) |
|
|
update_btn.click(fn=update_prompt, inputs=new_prompt, outputs=success_msg) |
|
|
|
|
|
demo.launch() |
|
|
|
