chatbot.py

import pandas as pd
import numpy as np
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
from spellchecker import SpellChecker
import re
import os
import logging
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification, Trainer, TrainingArguments
from torch.utils.data import Dataset
import torch
from sklearn.preprocessing import LabelEncoder

# Setup logging
logging.basicConfig(filename='chatbot.log', level=logging.INFO, format='%(asctime)s - %(message)s')

# Configuration
INTENTS = [
    "visa_application", "other", "travel", "job", "education", "culture",
    "integration", "healthcare", "market", "housing", "citizenship"
]
SYNONYMS = {
    "st paddy's day": "st patrick's day",
    "lepp card": "leap card",
    "resdence": "residence",
    "job": "employment",
    "study permit": "student visa",
    "irish stuff": "irish culture",
    "gnib": "irp card",
    "contact": "reach",
    "INIS": "immigration service",
    "pub toasts": "pub culture",
    "homesick": "homesickness",
    "travel": "transport"
}
IRISH_TERMS = {
    "fáilte": "welcome",
    "sláinte": "cheers",
    "céilí": "dance",
    "bealtaine": "May festival",
    "lá fhéile pádraig": "St. Patrick’s Day",
    "craic": "fun",
    "bodhrán": "drum",
    "gaeilge": "Irish language"
}
MAPPINGS = {
    "how to leap card?": ("how do i get a leap card in ireland?", "travel"),
    "what is contact of INIS?": ("how do i contact inis?", "visa_application"),
    "what’s the contact for INIS office?": ("how do i contact inis?", "visa_application"),
    "what’s the deal with st paddy’s day?": ("what is st patricks day and how can i enjoy it?", "culture"),
    "how do i live in Ireland?": ("how do i apply for permanent residency in ireland?", "visa_application"),
    "tell me about irish stuff.": ("what are irish music and dance traditions?", "culture"),
    "i am feeling home sick in ireland": ("how do i handle homesickness in dublin?", "integration"),
    "how can i get a work visa if i’m from brazil?": ("how do i apply for a work permit from brazil?", "visa_application"),
    "how to explore ireland as a tourist on a budget?": ("what are budget travel tips for tourists?", "travel"),
    "hello, do you know about GIP": ("what is the purpose of the start-up entrepreneur programme (step)?", "visa_application"),
    "tell me more about those pub toasts!": ("how do i learn about irish pub culture?", "culture"),
    "what’s the process for a student visa in Ireland?": ("what are the steps to apply for a study visa in ireland?", "visa_application"),
    "how to travel in dublin?": ("how do i travel around dublin?", "travel"),
    "i’m homesick in dublin, any tips?": ("how do i handle homesickness in dublin?", "integration"),
    "good night": ("none", "other"),
    "bye": ("none", "other"),
    "i am so good": ("none", "other")
}

# Preprocessing
spell = SpellChecker()

def normalize_query(query):
    query = query.lower().strip()
    query = re.sub(r'[^\w\s?]', '', query)
    return query

def correct_typos(query):
    words = query.split()
    corrected = [spell.correction(word) if spell.correction(word) else word for word in words]
    return " ".join(corrected)

def expand_synonyms(query):
    for synonym, canonical in SYNONYMS.items():
        query = query.replace(synonym, canonical)
    return query

def clean_irish_terms(response):
    response = response.replace("Antwort: ", "").replace("answer: ", "").strip(": ")
    for term, meaning in IRISH_TERMS.items():
        response = re.sub(rf"\*{re.escape(term)}\*\s*\([^)]*\)", f"*{term}*", response)
        response = re.sub(rf"\*{re.escape(term)}\*", f"*{term}* ({meaning})", response)
    return response

def preprocess_query(query):
    query = normalize_query(query)
    query = correct_typos(query)
    query = expand_synonyms(query)
    return query

# Intent Classification
class IntentDataset(Dataset):
    def __init__(self, questions, intents, tokenizer, max_length=128):
        self.questions = questions
        self.intents = intents
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.questions)

    def __getitem__(self, idx):
        question = self.questions[idx]
        intent = self.intents[idx]
        encoding = self.tokenizer(
            question,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].squeeze(),
            'attention_mask': encoding['attention_mask'].squeeze(),
            'labels': torch.tensor(intent, dtype=torch.long)
        }

class IntentClassifier:
    def __init__(self):
        self.tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
        self.model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=len(INTENTS))
        self.label_encoder = LabelEncoder()
        self.model_path = './distilbert_finetuned'

    def fine_tune(self, dataset):
        intents = self.label_encoder.fit_transform(dataset["intent"])
        questions = dataset["question"].tolist()
        df = pd.DataFrame({"question": questions, "intent": intents})
        balanced_df = []
        for intent in set(intents):
            intent_df = df[df["intent"] == intent]
            if len(intent_df) < 20:
                intent_df = pd.concat([intent_df] * (20 // len(intent_df) + 1), ignore_index=True)[:20]
            balanced_df.append(intent_df)
        balanced_df = pd.concat(balanced_df)
        train_dataset = IntentDataset(balanced_df["question"].tolist(), balanced_df["intent"].tolist(), self.tokenizer)
        training_args = TrainingArguments(
            output_dir=self.model_path,
            num_train_epochs=4,
            per_device_train_batch_size=8,
            warmup_steps=500,
            weight_decay=0.01,
            logging_dir='./logs',
            logging_steps=10,
            save_strategy="epoch"
        )
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=train_dataset
        )
        trainer.train()
        self.model.save_pretrained(self.model_path)
        self.tokenizer.save_pretrained(self.model_path)

    def predict(self, query):
        if os.path.exists(self.model_path):
            self.model = DistilBertForSequenceClassification.from_pretrained(self.model_path)
        inputs = self.tokenizer(query, return_tensors="pt", truncation=True, padding=True)
        outputs = self.model(**inputs)
        predicted_class = torch.argmax(outputs.logits, dim=1).item()
        return self.label_encoder.inverse_transform([predicted_class])[0]

# Retrieval
class Retriever:
    def __init__(self, dataset):
        self.dataset = dataset
        self.model = SentenceTransformer('all-MiniLM-L6-v2')
        self.tokenized_corpus = [q.lower().split() for q in dataset["question"]]
        self.bm25 = BM25Okapi(self.tokenized_corpus)
        self.question_embeddings = self.model.encode(dataset["question"].tolist())

    def check_mapping(self, query):
        if query.lower() in MAPPINGS:
            mapped_q, intent = MAPPINGS[query.lower()]
            return mapped_q, 0.9, intent
        return None, 0.0, None

    def handle_vague_query(self, query):
        if any(keyword in query.lower() for keyword in ["stuff", "things", "general"]) or \
           ("culture" in query.lower() and "irish" in query.lower()):
            culture_questions = self.dataset[self.dataset["intent"] == "culture"]
            if culture_questions.empty:
                return None, 0.0, None
            culture_indices = culture_questions.index.tolist()
            scores = self.bm25.get_scores(query.lower().split())
            culture_scores = scores[culture_indices]
            if len(culture_scores) == 0:
                return None, 0.0, None
            top_idx = culture_scores.argmax()
            return culture_questions["question"].iloc[top_idx], 0.9, "culture"
        return None, 0.0, None

    def combine_scores(self, bm25_scores, sent_scores, bm25_weight=0.4, sent_weight=0.6):
        bm25_norm = bm25_scores / (bm25_scores.max() + 1e-10)
        sent_norm = sent_scores / (sent_scores.max() + 1e-10)
        return bm25_weight * bm25_norm + sent_weight * sent_norm

    def retrieve(self, query, threshold=0.5):
        query = preprocess_query(query)

        # Check mappings
        mapped_q, score, intent = self.check_mapping(query)
        if mapped_q:
            return mapped_q, score, intent

        # Check vague queries
        vague_q, vague_score, vague_intent = self.handle_vague_query(query)
        if vague_q:
            return vague_q, vague_score, vague_intent

        # BM25 + SentenceTransformer
        tokenized_query = query.lower().split()
        bm25_scores = self.bm25.get_scores(tokenized_query)
        query_embedding = self.model.encode([query])
        sent_scores = cosine_similarity(query_embedding, self.question_embeddings)[0]
        combined_scores = self.combine_scores(bm25_scores, sent_scores)

        top_idx = combined_scores.argmax()
        similarity = combined_scores[top_idx]

        if similarity >= threshold:
            return self.dataset["question"].iloc[top_idx], similarity, self.dataset["intent"].iloc[top_idx]
        return None, 0.0, None

# Response Generation
class ResponseGenerator:
    def __init__(self):
        self.tokenizer = None
        self.model = None

    def generate(self, query, retrieved_q, dataset, intent, similarity):
        if not retrieved_q or retrieved_q == "none":
            if query.lower() in ["good night", "bye"]:
                return "Cheers, catch you later! *Sláinte* (cheers)!", []
            if query.lower() == "i am so good":
                return "Glad you’re feeling the *craic* (fun)! What’s up next for your Irish adventure?", []
            return f"Sorry, I couldn’t find a match! Try asking about visas, jobs, or Irish culture for some *craic* (fun) answers!", []
        answer = dataset[dataset["question"] == retrieved_q]["answer"].iloc[0]
        prefix = "Yo, here’s the *craic* (fun) on that: " if intent in ["culture", "travel"] else ""
        response = clean_irish_terms(prefix + answer.strip())
        return response, [(retrieved_q, similarity, intent)]

# Main Function
def main():
    # Load and clean dataset
    dataset = pd.read_csv("dataset.csv")
    dataset["question"] = dataset["question"].str.lower().str.replace(r"[^\w\s?]", "", regex=True)
    dataset["answer"] = dataset["answer"].str.replace(r"\*{2}([^*]+)\*{2}", r"*\1*", regex=True)
    dataset["intent"] = dataset["intent"].str.lower().replace("greeting", "other")
    print("Dataset Intent Distribution:")
    print(dataset["intent"].value_counts())

    # Initialize components
    retriever = Retriever(dataset)
    classifier = IntentClassifier()
    generator = ResponseGenerator()

    # Fine-tune classifier
    classifier.fine_tune(dataset)

    # Test queries
    test_queries = [
        "how to leap card?",
        "What’s the deal with St. Paddy’s Day?",
        "Tell me about Irish stuff.",
        "I am EU citizen do I need visa to work in Ireland?",
        "what is contact of INIS?",
        "how to get irp card?",
        "What is saint patrick’s day?",
        "I’m from Nigeria, how do I study in Ireland?",
        "hi",
        "Hi, How do I apply for a Work Permit?",
        "Tell me about artificial intelligence.",
        "What’s the weather like?",
        "how are you?",
        "Hello, I need help with my application.",
        "Hello, do I need work permit to work in Ireland?",
        "Hi, what is the process for getting a visa?",
        "Hello, How do I apply for a Study Permit in Ireland?",
        "Hi, Can I retire in Ireland?",
        "Hello, do you know about GIP",
        "good morning",
        "thanks for your answer",
        "I am from india do I need visa to work in Ireland?",
        "Good night",
        "Bye",
        "Thank you for your help!",
        "I am so good",
        "Can you tell some irish culture?",
        "what is farmer's market?",
        "How to attend a Gaelic football match?",
        "what is irp card",
        "how to get irish citizenship?",
        "how to get irish passport?",
        "how to travel in dublin?",
        "how to get a job in ireland?",
        "how apply for irish residency?",
        "How can I get a work visa if I’m from Brazil?",
        "What’s the process for a student visa in Ireland?",
        "I’m an EU citizen, do I need a permit to work in Dublin?",
        "How to obtain a residence card in Ireland?",
        "What’s a GNIB card?",
        "Where can I find cheap accommodation in Galway?",
        "How to rent an apartment in Ireland as a non-EU student?",
        "What’s the cost of living in Cork for a worker?",
        "What’s a *céilí* dance and how do I join one?",
        "How can I celebrate Bealtaine in Ireland?",
        "Tell me about Irish music with a *bodhrán*.",
        "Hey, what’s that festival with parades in March?",
        "Tell me more about those pub toasts!",
        "How to get an irish resdence card?",
        "What’s the contact for INIS office?",
        "Can I retire in Cork as a US citizen?",
        "I’m Australian, how do I get a job in Ireland?",
        "What’s the process for getting a PPS number?",
        "I’m homesick in Dublin, any tips?",
        "How to join Irish community groups as an immigrant?",
        "What’s a permit?",
        "How do I live in Ireland?"
    ]

    results = []
    for query in test_queries:
        logging.info(f"Processing query: {query}")
        preprocessed_query = preprocess_query(query)
        logging.info(f"Preprocessed query: {preprocessed_query}")

        retrieved_q, similarity, intent = retriever.retrieve(query)
        logging.info(f"Retrieved: {retrieved_q}, Similarity: {similarity}, Intent: {intent}")

        if not retrieved_q:
            intent = classifier.predict(preprocessed_query)
            logging.info(f"Fallback intent: {intent}")

        response, top_matches = generator.generate(query, retrieved_q, dataset, intent, similarity)
        logging.info(f"Response: {response}")

        result = {
            "Query": query,
            "Retrieved Question": retrieved_q,
            "Similarity": similarity,
            "Predicted Intent": intent,
            "Response": response,
            "Top Matches": top_matches
        }
        results.append(result)
        print(f"Query: {result['Query']}")
        print(f"Debug: Query={result['Query']}, Similarity={result['Similarity']:.3f}, Predicted Intent={result['Predicted Intent']}, Retrieved Q={result['Retrieved Question']}")
        print(f"Response: {result['Response']}")
        print(f"Top Matches: {result['Top Matches']}")
        print()

    # Save results
    with open("chatbot_results.txt", "w") as f:
        for result in results:
            f.write(f"Query: {result['Query']}\n")
            f.write(f"Debug: Query={result['Query']}, Similarity={result['Similarity']:.3f}, Predicted Intent={result['Predicted Intent']}, Retrieved Q={result['Retrieved Question']}\n")
            f.write(f"Response: {result['Response']}\n")
            f.write(f"Top Matches: {result['Top Matches']}\n\n")

if __name__ == "__main__":
    main()