Spaces:
Sleeping
Sleeping
File size: 4,256 Bytes
a2f0146 6ea7a6e a8452a4 a2f0146 2a46434 a2f0146 a8452a4 a2f0146 be72f4b a2f0146 a8452a4 a2f0146 a8452a4 be72f4b a2f0146 2a46434 a2f0146 a8452a4 a2f0146 a8452a4 a2f0146 a8452a4 a2f0146 be72f4b a2f0146 2a46434 a2f0146 66990f5 a2f0146 a8452a4 a2f0146 66990f5 a2f0146 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 |
import time
import gradio as gr
from datasets import load_dataset, Dataset
from huggingface_hub import hf_hub_download
from sentence_transformers import SentenceTransformer, util
import torch
# ---------------------------
# CONFIGURATION
# ---------------------------
HF_TOKEN = "<YOUR_HF_TOKEN>" # set your HF token
DATASET_NAME = "guardian-ai-qna"
MAX_QUESTIONS = 5 # max questions per TIME_WINDOW
TIME_WINDOW = 3600 # 1 hour in seconds
EMBED_MODEL = "all-MiniLM-L6-v2" # small but effective embedding model
# ---------------------------
# LOAD OR CREATE DATASET
# ---------------------------
try:
dataset = load_dataset(DATASET_NAME, use_auth_token=HF_TOKEN)
dataset = dataset["train"]
except:
dataset = Dataset.from_dict({"question": [], "answer": []})
# ---------------------------
# EMBEDDING MODEL
# ---------------------------
embedder = SentenceTransformer(EMBED_MODEL)
# Precompute embeddings for existing Q&A
if len(dataset) > 0:
dataset_embeddings = embedder.encode(dataset["question"], convert_to_tensor=True)
else:
dataset_embeddings = torch.empty((0, embedder.get_sentence_embedding_dimension()))
# ---------------------------
# USER RATE LIMITING
# ---------------------------
user_limits = {}
def check_rate_limit(session_id):
current_time = time.time()
if session_id not in user_limits:
user_limits[session_id] = {"count": 0, "start_time": current_time}
user_data = user_limits[session_id]
if current_time - user_data["start_time"] > TIME_WINDOW:
user_data["count"] = 0
user_data["start_time"] = current_time
if user_data["count"] >= MAX_QUESTIONS:
return False, f"You have reached the max of {MAX_QUESTIONS} questions. Please wait before asking more."
user_data["count"] += 1
return True, None
# ---------------------------
# HELPER FUNCTIONS
# ---------------------------
def find_similar_answer(user_input):
if len(dataset) == 0:
return None
query_emb = embedder.encode(user_input, convert_to_tensor=True)
scores = util.cos_sim(query_emb, dataset_embeddings)
top_idx = torch.argmax(scores)
top_score = scores[0][top_idx].item()
if top_score > 0.6: # threshold for similarity
return dataset["answer"][top_idx]
return None
def save_qna(question, answer):
global dataset, dataset_embeddings
new_entry = Dataset.from_dict({"question": [question], "answer": [answer]})
dataset = Dataset.from_dict({
"question": dataset["question"] + new_entry["question"],
"answer": dataset["answer"] + new_entry["answer"]
})
# update embeddings incrementally
new_emb = embedder.encode([question], convert_to_tensor=True)
if len(dataset_embeddings) == 0:
dataset_embeddings = new_emb
else:
dataset_embeddings = torch.vstack([dataset_embeddings, new_emb])
# save to HF dataset (push to hub)
dataset.push_to_hub(DATASET_NAME, token=HF_TOKEN)
# ---------------------------
# MAIN CHAT FUNCTION
# ---------------------------
def chat(history, user_input, session_id="default"):
# Rate limit check
allowed, message = check_rate_limit(session_id)
if not allowed:
history.append(("System", message))
return history, history
# Check existing similar Q&A
response = find_similar_answer(user_input)
if not response:
# Fallback / simple generative response
response = f"Guardian AI: Sorry, I don’t know the answer yet. I’m learning!"
# Save new Q&A for incremental learning
save_qna(user_input, response)
# Update chat history
history.append((user_input, response))
return history, history
# ---------------------------
# GRADIO INTERFACE
# ---------------------------
with gr.Blocks() as app:
chatbot = gr.Chatbot()
msg = gr.Textbox(label="Your question")
session_state = gr.State("default") # default session
def user_submit(message, history, session_id):
return chat(history, message, session_id)
msg.submit(user_submit, inputs=[msg, chatbot, session_state], outputs=[chatbot, chatbot])
# Launch app
app.launch(server_name="0.0.0.0", server_port=7860, share=True)
|