from transformers import AutoModelForCausalLM, AutoTokenizer, MarianMTModel, MarianTokenizer
import torch
import re
import json

MODEL_ID = "HuggingFaceTB/SmolLM2-360M"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32, low_cpu_mem_usage=True)

TRANS_MODEL_ID = "Helsinki-NLP/opus-mt-en-zh"
trans_tokenizer = MarianTokenizer.from_pretrained(TRANS_MODEL_ID)
trans_model = MarianMTModel.from_pretrained(TRANS_MODEL_ID)

def clean_output(text: str) -> str:
    text = re.sub(r"</?[a-zA-Z0-9]+[^>]*>", "", text)
    text = text.replace("<|endoftext|>", "").replace("<|im_end|>", "").replace("<|im_start|>", "")
    text = re.sub(r"(?m)^\s*(?:[0-9]+[.\、\)]|[一二三四五六七八九十]+[、\.]|[（\(][0-9一二三四五六七八九十]+[）\)]|[①②③④⑤⑥⑦⑧⑨⑩]|(?:第[一二三四五六七八九十0-9]+[条点个部分阶段、：:]))\s*", "", text)
    text = re.sub(r"\s+[0-9]+[.\、]\s*", " ", text)
    text = re.sub(r"\s+[一二三四五六七八九十]+[、]\s*", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    match = re.search(r"[。！？\n][^。！？\n]*$", text)
    if match and match.start() > 15:
        text = text[:match.start() + 1]
    return text.strip()

def translate_to_chinese(text: str) -> str:
    if not text or len(text.strip()) < 3:
        return text
    inputs = trans_tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    with torch.no_grad():
        translated = trans_model.generate(**inputs, max_new_tokens=256)
    return trans_tokenizer.decode(translated[0], skip_special_tokens=True)

def clean_chinese(text: str) -> str:
    cleaned = re.sub(r'[a-zA-Z\'\-]{3,}', '', text)
    cleaned = re.sub(r'[^\u4e00-\u9fff\u3000-\u303f\uff00-\uffef，。！？、；：""''（）—…\s]', '', cleaned)
    cleaned = re.sub(r'\s{2,}', ' ', cleaned)
    cleaned = cleaned.strip()
    cleaned = re.sub(r'，\s*(并且|而且|但是|可是|然而|而|因此|所以|于是|因为|由于|虽然)', r'。\1', cleaned)
    segments = cleaned.split('，')
    result = []
    current_len = 0
    for i, seg in enumerate(segments):
        result.append(seg)
        current_len += len(seg)
        if i < len(segments) - 1:
            if current_len > 12:
                result.append('。')
                current_len = 0
            else:
                result.append('，')
    cleaned = "".join(result)
    cleaned = re.sub(r'。+', '。', cleaned)
    cleaned = cleaned.replace('。，', '。').replace('，。', '。')
    if cleaned and cleaned[-1] not in ['。', '！', '？', '”']:
        cleaned += '。'
    if len(cleaned) < 5:
        return ""
    return cleaned

prompt = "Finally, remember that everyone possesses unique gifts worth recognizing—not just numerical rankings or social status. Embrace diversity, celebrate individuality, cherish relationships built on trust and mutual respect, and stay curious about ever-expanding horizons as we did here today. Happy adventures!"
inputs = tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
    output_ids = model.generate(
        **inputs, max_new_tokens=120, do_sample=True, temperature=1.05, top_p=0.92, repetition_penalty=1.15
    )
generated_tokens = output_ids[0][inputs.input_ids.shape[1]:]
raw_content = tokenizer.decode(generated_tokens, skip_special_tokens=True)
raw_content = clean_output(raw_content)
translated = translate_to_chinese(raw_content)
cleaned = clean_chinese(translated)

print("raw:", raw_content)
print("trans:", translated)
print("clean:", cleaned)