import jsonimport reimport torchfrom transformers import AutoModelForCausalLM, AutoTokenizer MODEL_ID = "Salahin/youtube-audiobook-title-parser" tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=torch.float16, device_map="auto")model.eval() def parse_title(title: str) -> dict: messages = [ { "role": "system", "content": "Extract book title, author, and channel as JSON. Keep Bengali in Bengali script." }, {"role": "user", "content": title} ] prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer([prompt], return_tensors="pt").to(model.device) with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=80, do_sample=False, eos_token_id=tokenizer.convert_tokens_to_ids("<|im_end|>") ) text = tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True).strip() match = re.search(r'\{.*?\}', text, re.DOTALL) return json.loads(match.group(0)) if match else {"raw": text} # Exampleprint(parse_title("হিমুর হাতে কয়েকটি নীলপদ্ম 1/2 | Humayun Ahmed | Golpokothon"))# Output: {'book_title': 'হিমুর হাতে কয়েকটি নীলপদ্ম', 'author': 'Humayun Ahmed', 'channel': 'Golpokothon'}