Spaces:
Sleeping
Sleeping
| import torch | |
| from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline | |
| from sentence_transformers import SentenceTransformer, util | |
| import gradio as gr | |
| import json | |
| # Load the lightweight BERT-based QA model optimized for CPU | |
| model_name = "distilbert-base-uncased-distilled-squad" # Efficient for CPU | |
| model = AutoModelForQuestionAnswering.from_pretrained(model_name) | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| # Initialize pipeline for CPU usage | |
| device = -1 # Force CPU | |
| qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer, device=device) | |
| # Load Sentence-BERT for semantic search | |
| embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2') | |
| # Load knowledge base and expanded QA dataset | |
| with open('knowledge_base.json', 'r') as f: | |
| knowledge_base = json.load(f) | |
| with open('expanded_qa_dataset.json', 'r') as f: | |
| expanded_qa_dataset = json.load(f) | |
| # Function to create embeddings for the expanded QA dataset | |
| def create_qa_dataset_embeddings(expanded_qa_dataset): | |
| qa_embeddings = [] | |
| questions = [] | |
| for item in expanded_qa_dataset: | |
| questions.append(item['question']) | |
| qa_embeddings.append(embedding_model.encode(item['question'], convert_to_tensor=True)) | |
| return qa_embeddings, questions | |
| # Create QA dataset embeddings | |
| qa_embeddings, qa_questions = create_qa_dataset_embeddings(expanded_qa_dataset) | |
| # Function to create embeddings for the knowledge base content | |
| def create_knowledge_base_embeddings(knowledge_base): | |
| embeddings = [] | |
| for entry in knowledge_base: | |
| if 'title' in entry: | |
| content = entry['title'] + ' '.join( | |
| [c.get('text', '') for c in entry.get('content', [])] + | |
| [' '.join(step['details']) for c in entry.get('content', []) if 'steps' in c for step in c['steps']] + | |
| [faq['question'] + ' ' + faq['answer'] for c in entry.get('content', []) if 'faq' in c for faq in c['faq']] | |
| ) | |
| embeddings.append(embedding_model.encode(content, convert_to_tensor=True)) | |
| return embeddings | |
| # Create knowledge base embeddings | |
| knowledge_base_embeddings = create_knowledge_base_embeddings(knowledge_base) | |
| # Semantic search on expanded QA dataset | |
| def search_expanded_qa(question): | |
| question_embedding = embedding_model.encode(question, convert_to_tensor=True) | |
| cosine_scores = util.pytorch_cos_sim(question_embedding, torch.stack(qa_embeddings)) | |
| best_match_idx = torch.argmax(cosine_scores).item() | |
| best_match_score = cosine_scores[0, best_match_idx].item() | |
| return expanded_qa_dataset[best_match_idx]['answer'], best_match_score | |
| # Semantic search on knowledge base | |
| def search_knowledge_base(question): | |
| question_embedding = embedding_model.encode(question, convert_to_tensor=True) | |
| cosine_scores = util.pytorch_cos_sim(question_embedding, torch.stack(knowledge_base_embeddings)) | |
| best_match_idx = torch.argmax(cosine_scores).item() | |
| best_match_score = cosine_scores[0, best_match_idx].item() | |
| # Retrieve content from best matched knowledge base entry | |
| best_match_entry = knowledge_base[best_match_idx] | |
| for content_item in best_match_entry['content']: | |
| if 'faq' in content_item: | |
| for faq in content_item['faq']: | |
| if faq['question'].lower() in question.lower(): | |
| return faq['answer'], best_match_score | |
| if 'steps' in content_item: | |
| step_details = [step['details'] for step in content_item['steps']] | |
| return "\n".join(step_details), best_match_score | |
| if 'text' in content_item: | |
| return content_item['text'], best_match_score | |
| return "Lo siento, no encontré una respuesta adecuada para tu pregunta.", best_match_score | |
| # Answer function: search both datasets and return the best match | |
| def answer_question(question): | |
| # Search expanded QA dataset | |
| qa_answer, qa_score = search_expanded_qa(question) | |
| # Search knowledge base | |
| kb_answer, kb_score = search_knowledge_base(question) | |
| # Compare scores and return the best answer | |
| if qa_score >= kb_score: | |
| return qa_answer | |
| else: | |
| return kb_answer | |
| # Gradio interface | |
| interface = gr.Interface( | |
| fn=answer_question, | |
| inputs="text", | |
| outputs="text", | |
| title="OCN Customer Support Chatbot", | |
| description="Ask questions and get answers from the OCN knowledge base and expanded QA dataset." | |
| ) | |
| # Launch the interface | |
| interface.launch(share=True) | |