Diabetrek_AI / chatbot.py
EdgarDataScientist's picture
chatbot.py
a5dd441 verified
import torch
import fitz # PyMuPDF for PDF text extraction
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from sentence_transformers import SentenceTransformer
import faiss
import gradio as gr
import os
from huggingface_hub import login
# Authenticate with Hugging Face Hub
pdf_path1 ='/content/Chrono 1.pdf'
pdf_path2 ='/content/Chrono 2.pdf'
pdf_path3 ='/content/Chrono 3.pdf'
# Load the Mistral 7B model and tokenizer
model_name = 'mistralai/Mistral-7B-Instruct-v0.3'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
# Load sentence transformer for embedding and similarity search
embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
# Function to extract text from PDFs
def extract_text_from_pdf(pdf_file_path):
doc = fitz.open(pdf_file_path)
text = ""
for page in doc:
text += page.get_text("text")
return text
# Placeholder PDF knowledge base (Extracted content from PDFs)
pdf_knowledge_base = []
pdf_files = [pdf_path1, pdf_path2,pdf_path3] # Add the actual paths to your PDF files
for pdf_file in pdf_files:
pdf_text = extract_text_from_pdf(pdf_file)
pdf_knowledge_base.append({"document": pdf_file, "content": pdf_text})
# Combine extracted text with specific company information
knowledge_base = [
{"question": "How does DiabeTrek ensure data privacy and security?",
"answer": ("DiabeTrek ensures data privacy through multiple layers of protection, including data encryption during "
"transit and at rest. We comply with regulations like HIPAA and GDPR to safeguard your personal data.")},
{"question": "What are DiabeTrek's emergency guidelines?",
"answer": "DiabeTrek advises users to seek immediate medical attention in case of diabetes-related emergencies. This chatbot is not for emergency use."},
{"question": "What are DiabeTrek's mission, vision, and values?",
"answer": "DiabeTrek's mission is to improve the lives of people with diabetes through innovative AI-driven solutions. Our vision is a world where diabetes care is seamless, proactive, and accessible."},
# Additional items can be added here following the CEO's instructions
]
# Create a FAISS index for efficient retrieval
embedding_dim = 384 # Output dimension of the MiniLM model
index = faiss.IndexFlatL2(embedding_dim)
# Create a list of embeddings and index them
knowledge_embeddings = []
for entry in knowledge_base:
embedding = embedder.encode(entry['question'], convert_to_tensor=False)
knowledge_embeddings.append(embedding)
index.add(embedding.reshape(1, -1))
# Create embeddings for PDF content and index them
for pdf_entry in pdf_knowledge_base:
embedding = embedder.encode(pdf_entry['content'], convert_to_tensor=False)
knowledge_embeddings.append(embedding)
index.add(embedding.reshape(1, -1))
# RAG Retrieval function
def retrieve_knowledge(question, top_k=1):
question_embedding = embedder.encode(question, convert_to_tensor=False)
D, I = index.search(question_embedding.reshape(1, -1), top_k)
results = [knowledge_base[idx] for idx in I[0]]
return results
# Chatbot function combining retrieval and generation
def customer_support_chatbot(user_input):
# Retrieve relevant knowledge
retrieved_knowledge = retrieve_knowledge(user_input)
# Prepare context for the generative model
context = " ".join([f"Q: {entry['question']} A: {entry['answer']}" for entry in retrieved_knowledge])
# Generate response using Mistral
prompt = f"Customer Question: {user_input}\n\n{context}\n\nResponse:"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_length=150, do_sample=True, temperature=0.7)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response
# Gradio UI
def gradio_interface(user_input):
response = customer_support_chatbot(user_input)
return response
# Build Gradio interface
interface = gr.Interface(fn=gradio_interface,
inputs="text",
outputs="text",
title="DiabeTrek Customer Support Chatbot",
description="Ask any question about DiabeTrek, its services, and policies.")
# Launch the Gradio app
interface.launch()