# Somali Chatbot with Translation + English Chatbot using Gradio

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline, AutoModelForCausalLM
import torch

# Load Translation Model: facebook/nllb-200-3.3B
trans_model = "facebook/nllb-200-3.3B"
trans_tokenizer = AutoTokenizer.from_pretrained(trans_model)
trans_model_loaded = AutoModelForSeq2SeqLM.from_pretrained(trans_model)

def translate(text, src_lang, tgt_lang):
    translator = pipeline("translation",
                          model=trans_model_loaded,
                          tokenizer=trans_tokenizer,
                          src_lang=src_lang,
                          tgt_lang=tgt_lang,
                          max_length=512)
    return translator(text)[0]['translation_text']

# Load Chatbot Model: meta-llama/Meta-Llama-3-8B-Instruct
chatbot_model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
chatbot_tokenizer = AutoTokenizer.from_pretrained(chatbot_model_name)
chatbot_model = AutoModelForCausalLM.from_pretrained(
    chatbot_model_name, torch_dtype=torch.float16, device_map="auto")

def ask_chatbot(prompt):
    input_ids = chatbot_tokenizer.encode(prompt, return_tensors="pt").to(chatbot_model.device)
    with torch.no_grad():
        output = chatbot_model.generate(input_ids, max_new_tokens=256, do_sample=True, temperature=0.7)
    return chatbot_tokenizer.decode(output[0], skip_special_tokens=True)

# Full pipeline function
def somali_chatbot_pipeline(somali_input):
    # Step 1: Somali -> English
    english_input = translate(somali_input, src_lang="som_Latn", tgt_lang="eng_Latn")

    # Step 2: Ask English chatbot
    chatbot_response = ask_chatbot(english_input)

    # Step 3: English -> Somali
    somali_output = translate(chatbot_response, src_lang="eng_Latn", tgt_lang="som_Latn")

    return somali_output

# Gradio Interface (with submit button)
def build_interface():
    with gr.Blocks() as demo:
        gr.Markdown("""
        # 🤖 Chatbot Soomaaliyeed 
        Qor su'aal Af Soomaali ah. Nidaamku wuxuu si qarsoodi ah ugu rogaa luqadda Ingiriisiga, ka dib wuxuu weydiiyaa chatbot English ah, kadibna jawaabta wuxuu kuu soo celinayaa Af Soomaali.
        """)

        with gr.Row():
            somali_input = gr.Textbox(label="Su’aashaada Af Soomaali", placeholder="Maxaa keena macaanka?", lines=3)
        with gr.Row():
            submit_btn = gr.Button("Dir Su’aasha")
        with gr.Row():
            somali_output = gr.Textbox(label="Jawaabta Chatbot-ka (Af Soomaali)", lines=5)

        submit_btn.click(fn=somali_chatbot_pipeline, inputs=somali_input, outputs=somali_output)

    return demo

if __name__ == "__main__":
    iface = build_interface()
    iface.launch()