phpcoder/zimeng-chat

1 2import gradio as gr3from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer4import torch5from typing import List, Dict, Tuple6import json7import os8from datetime import datetime9 10 11CONFIG = {12    "allowed_topics": [13        "government services", "community issues", "public services",14        "zimbabwe government", "local authorities", "public utilities",15        "infrastructure", "health services", "education services",16        "transportation", "waste management", "public safety"17    ],18    "min_confidence": 0.5,19    "default_response": "I can only answer questions about government and community services in Zimbabwe. How can I assist you with that?",20    "max_history": 5  21}22 23# Load models24def load_models():25    26    intent_classifier = pipeline(27        "zero-shot-classification",28        model="facebook/bart-large-mnli",29        device=0 if torch.cuda.is_available() else -130    )31    translator = pipeline(32        "translation_en_to_xx",33        model="Helsinki-NLP/opus-mt-en-romance",  34        device=0 if torch.cuda.is_available() else -135    )36    37 38    local_llm = pipeline(39        "text-generation",40        model="gpt2", 41        device=0 if torch.cuda.is_available() else -142    )43    44    return {45        "intent_classifier": intent_classifier,46        "translator": translator,47        "local_llm": local_llm48    }49 50# Initialize models51models = load_models()52 53def detect_language(text: str) -> str:54    """Detect language using simple heuristic (can be replaced with langdetect)"""55    common_zw_languages = {56        'en': set(['the', 'and', 'you', 'that', 'have']),57        'sn': set(['ne', 'uye', 'uye', 'kuti', 'uye']),  # Shona58        'nd': set(['kuti', 'uye', 'uye', 'uye', 'uye'])   # Ndebele59    }60    61    words = set(text.lower().split()[:10])  62    scores = {63        lang: len(words.intersection(vocab)) 64        for lang, vocab in common_zw_languages.items()65    }66    return max(scores.items(), key=lambda x: x[1])[0] if scores else 'en'67 68def translate_text(text: str, target_lang: str = 'en') -> str:69    """Translate text to target language"""70    if target_lang == 'en':71        return text72        73    try:74        if target_lang in ['sn', 'nd']:  75            translation = models['translator'](76                text,77                src_lang="eng_Latn",78                tgt_lang="sna_Latn" if target_lang == 'sn' else "nde_Latn"79            )80            return translation[0]['translation_text']81        return text82    except Exception as e:83        print(f"Translation error: {e}")84        return text85 86def is_query_allowed(text: str) -> Tuple[bool, float]:87    """Check if the query is within allowed topics"""88    try:89        result = models['intent_classifier'](90            text,91            candidate_labels=CONFIG["allowed_topics"],92            multi_label=True93        )94        max_confidence = max(result['scores']) if result['scores'] else 095        return max_confidence >= CONFIG["min_confidence"], max_confidence96    except Exception as e:97        print(f"Intent classification error: {e}")98        return False, 099 100def generate_response(prompt: str, chat_history: List[Tuple[str, str]] = None) -> str:101    """Generate a response using the local model"""102    try:103        context = "\n".join([f"User: {msg[0]}\nBot: {msg[1]}" for msg in (chat_history or [])[-CONFIG["max_history"]:]])104        full_prompt = f"{context}\nUser: {prompt}\nBot:"105 106        response = models['local_llm'](107            full_prompt,108            max_length=150,109            num_return_sequences=1,110            temperature=0.7,111            top_p=0.9,112            do_sample=True113        )114        115        if response and len(response) > 0:116            return response[0]['generated_text'].split("Bot:")[-1].strip()117        return "I'm not sure how to respond to that."118    except Exception as e:119        print(f"Response generation error: {e}")120        return "I encountered an error processing your request."121 122def chat_interface(message: str, history: List[Tuple[str, str]]) -> str:123    """Main chat interface function"""124  125    lang = detect_language(message)126    127    if lang != 'en':128        message_en = translate_text(message, 'en')129    else:130        message_en = message131    132    # Check if query is allowed133    is_allowed, confidence = is_query_allowed(message_en)134    135    if not is_allowed:136        return CONFIG["default_response"]137    138    # Generate response139    response = generate_response(message_en, history)140    141    142    if lang != 'en':143        response = translate_text(response, lang)144    145    return response146 147 148with gr.Blocks(title="ZimEngage Chatbot") as demo:149    gr.Markdown("# ZimEngage Government Services Assistant")150    gr.Markdown("Ask me about government services and community issues in Zimbabwe.")151    152    chatbot = gr.Chatbot()153    msg = gr.Textbox(label="Your Message", placeholder="Type your message here...")154    clear = gr.Button("Clear")155    156    def respond(message, chat_history):157        bot_message = chat_interface(message, chat_history)158        chat_history.append((message, bot_message))159        return "", chat_history160    161    msg.submit(respond, [msg, chatbot], [msg, chatbot])162    clear.click(lambda: None, None, chatbot, queue=False)163 164# Run the app165if __name__ == "__main__":166    demo.launch(debug=True)