import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel, PeftConfig
import spaces
import time

model_name = "Anassk/MoroccanDarija-Llama-3.1-8B"
base_model_name = "unsloth/Meta-Llama-3.1-8B"

# Load tokenizer globally
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

@spaces.GPU
def load_model():
    try:
        # Load the base model
        base_model = AutoModelForCausalLM.from_pretrained(
            base_model_name,
            torch_dtype=torch.float16,
            device_map="auto",
            trust_remote_code=True,
            low_cpu_mem_usage=True
        )
        # Load the PEFT model
        model = PeftModel.from_pretrained(base_model, model_name)
        print(f"Model loaded successfully. Using device: {model.device}")
        return model
    except Exception as e:
        print(f"Error loading model: {e}")
        raise

@spaces.GPU
def generate_text(prompt):
    model = load_model()
    unique_id = int(time.time() * 1000)  # Add a unique identifier
    full_prompt = f"Task {unique_id}: Translate this exact Darija text to English and explain its meaning. Do not add or change any words. Darija text: {prompt}"
    inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=512)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=200, num_return_sequences=1, temperature=0.7)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Remove the original prompt and any potential prefixes
    response = response.replace(full_prompt, "").strip()
    response = response.split("Task " + str(unique_id) + ":")[-1].strip()
    
    # Check if the response contains the original Darija text
    if prompt not in response:
        response = f"Error: Model response does not match the input. Original input: {prompt}\n\nModel output: {response}"
    
    # Format the output
    darija_text = f"Darija (Arabic script): {prompt}"
    english_translation = f"English translation and explanation: {response}"
    
    formatted_response = f"{darija_text}\n\n{english_translation}"
    return formatted_response

gradio_app = gr.Interface(
    generate_text,
    inputs=gr.Textbox(label="Enter your prompt in Darija", lines=3),
    outputs=gr.Textbox(label="Model Output"),
    title="Moroccan Darija Language Model",
    description=f"""This model is trained on Moroccan Darija and can translate Darija to English and provide explanations.
    It is fine-tuned from the {base_model_name} model.
    Enter your prompt in Darija and see how the model translates and explains it!""",
    examples=[
        ["واش كاين شي حد هنا كيهضر الدارجة؟"],
        ["غادي نمشي للسوق باش نشري الخضرة"],
        ["حكي لي على شي قصة مغربية قديمة"],
        ["شنو هي أحسن طريقة باش نتعلم الدارجة؟"],
        ["واش كاين شي مثل شعبي على الصبر؟"],
        ["كيفاش كنديرو الطاجين المغربي؟"],
        ["شنو هي العادات ديال رمضان فالمغرب؟"]
    ]
)

if __name__ == "__main__":
    gradio_app.launch()