| import os |
| import gradio as gr |
| import random as rand |
| import torch |
| from huggingface_hub import InferenceClient |
| from sentence_transformers import SentenceTransformer |
|
|
|
|
| custom_css = """ |
| @import url('https://fonts.googleapis.com/css2?family=Fredoka:wght@300..700&display=swap'); |
| body, input, textarea, button, select, label, p, span, div, h1, h2, h3, h4, h5, h6 { |
| font-family: 'Fredoka', sans-serif !important; |
| } |
| """ |
|
|
| with open("knowledgebase.txt", "r", encoding="utf-8") as file: |
| knowledge_base = file.read() |
|
|
| def preprocess_text(text): |
| cleaned_text = text.strip() |
| chunks = cleaned_text.split("\n") |
| cleaned_chunks = [] |
| for chunk in chunks: |
| stripped_chunk = chunk.strip() |
| if len(stripped_chunk) > 0: |
| cleaned_chunks.append(stripped_chunk) |
| return cleaned_chunks |
|
|
| cleaned_chunks = preprocess_text(knowledge_base) |
|
|
| model = SentenceTransformer('all-MiniLM-L6-v2') |
|
|
| def create_embeddings(text_chunks): |
| chunk_embeddings = model.encode(text_chunks, convert_to_tensor=True) |
| print("Embeddings Shape:", chunk_embeddings.shape) |
| return chunk_embeddings |
|
|
| chunk_embeddings = create_embeddings(cleaned_chunks) |
|
|
| def get_top_chunks(query, chunk_embeddings, text_chunks): |
| query_embedding = model.encode(query, convert_to_tensor=True) |
| query_embedding_normalized = query_embedding / query_embedding.norm() |
| chunk_embeddings_normalized = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True) |
|
|
| similarities = torch.matmul(chunk_embeddings_normalized, query_embedding_normalized) |
| top_indices = torch.topk(similarities, k=3).indices |
| |
| top_chunks = [] |
| for i in top_indices: |
| chunk = text_chunks[i] |
| top_chunks.append(chunk) |
| return top_chunks |
|
|
|
|
| client = InferenceClient("Qwen/Qwen2.5-7B-Instruct", token = os.getenv("hf")) |
|
|
| def respond(message, history): |
| relevant_context = get_top_chunks(message, chunk_embeddings, cleaned_chunks) |
| context_str = "\n".join(relevant_context) |
| |
| system_prompt = f"You are a helpful assistant. Use the following context to answer the user's question accurately:\n\n{context_str}" |
| |
| messages = [{"role": "system", "content": system_prompt}] |
| |
| if history: |
| messages.extend(history) |
| |
| messages.append({"role": "user", "content": message}) |
| |
| response = "" |
| for msg in client.chat_completion(messages, max_tokens=200, temperature=1, top_p=0.5, stream=True): |
| token = msg.choices[0].delta.content |
| if token: |
| response += token |
| yield response |
|
|
|
|
| with gr.Blocks(theme='d8ahazard/material_design_rd', css=custom_css) as demo: |
| gr.Markdown("# Digital Wellbeing") |
| |
| chatbot = gr.Chatbot(label="Chat History") |
| msg = gr.Textbox(placeholder="Type your question here...") |
| |
| with gr.Row(): |
| submit_btn = gr.Button("Submit", variant="primary") |
| clear_btn = gr.ClearButton([msg, chatbot], value="Clear Conversation") |
|
|
| submit_event = submit_btn.click(respond, inputs=[msg, chatbot], outputs=[chatbot]) |
| msg_submit_event = msg.submit(respond, inputs=[msg, chatbot], outputs=[chatbot]) |
| |
| submit_event.then(lambda: "", None, [msg]) |
| msg_submit_event.then(lambda: "", None, [msg]) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
