import spaces
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "At-Tawheed/Anis"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype=torch.bfloat16,
    device_map="auto"
)

SYSTEM_PROMPT = "You are ATTLAB, a helpful, harmless, and honest AI assistant developed by the ATTLAB team."


@spaces.GPU
def generate(message, history):
    try:
        def _extract_text(content):
            if isinstance(content, str):
                return content
            if isinstance(content, list):
                parts = []
                for item in content:
                    if isinstance(item, dict) and item.get("type") == "text":
                        parts.append(item.get("text", ""))
                    elif isinstance(item, str):
                        parts.append(item)
                return "".join(parts)
            return str(content)

        messages = [{"role": "system", "content": SYSTEM_PROMPT}]
        for turn in history:
            messages.append({"role": turn["role"], "content": _extract_text(turn["content"])})
        messages.append({"role": "user", "content": _extract_text(message)})

        inputs = tokenizer.apply_chat_template(
            messages,
            return_tensors="pt",
            add_generation_prompt=True,
            return_dict=True
        ).to(model.device)

        im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
        eos_ids = [tokenizer.eos_token_id]
        if im_end_id is not None and im_end_id != tokenizer.unk_token_id:
            eos_ids.append(im_end_id)

        output = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            do_sample=True,
            top_p=0.9,
            eos_token_id=eos_ids,
            pad_token_id=tokenizer.eos_token_id
        )
        response = tokenizer.decode(
            output[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True
        )
        return response if response else "(empty response)"
    except Exception as e:
        import traceback
        traceback.print_exc()
        return f"Error during generation: {e}"

demo = gr.ChatInterface(
    generate,
    title="Anis — ATTLAB",
    description="8B SFT model fine-tuned from Qwen2.5-7B by ATTLAB",
)

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