from peft import AutoPeftModelForCausalLM
from transformers import GenerationConfig, AutoTokenizer
import torch
import time


def generate_response(input_text: str) -> str:
    """
    Generate a response for the given input text using the Typhoon-7B model.

    Parameters:
    input_text (str): The input text prompt.

    Returns:
    str: The generated response.
    """
    # Initialize the tokenizer and model only once
    tokenizer = AutoTokenizer.from_pretrained("Thaweewat/typhoon-7b-chat-alpaca")

    model = AutoPeftModelForCausalLM.from_pretrained(
        "Thaweewat/typhoon-7b-chat-alpaca",
        low_cpu_mem_usage=True,
        return_dict=True,
        torch_dtype=torch.float16,
        device_map="cuda")

    generation_config = GenerationConfig(
        do_sample=True,
        top_k=1,
        temperature=0.5,
        max_new_tokens=300,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id)

    # Tokenize input
    inputs = tokenizer(input_text, return_tensors="pt").to("cuda")

    # Generate outputs
    st_time = time.time()
    outputs = model.generate(**inputs, generation_config=generation_config)

    # Decode and print response
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"Response time: {time.time() - st_time} seconds")
    return response

# Sample usage:
input_text = "###Human: ใครคือนายกไทยคนปัจจุบัน ###Assistant: "
print(generate_response(input_text))

"""
นายกรัฐมนตรีคนปัจจุบันของไทยคือพลเอกประยุทธ์ จันทร์โอชา เขาดำรงตำแหน่งนี้ตั้งแต่เดือนสิงหาคม 2557
และได้รับเลือกอีกครั้งในการเลือกตั้งทั่วไปในปี 2562 เขาเป็นนายทหารที่เกษียณอายุแล้วและเคยดำรงตำแหน่งผู้บัญชาการ
ทหารบกและผู้บัญชาการทหารสูงสุดมาก่อน เขาเป็นผู้นำรัฐบาลทหารตั้งแต่เดือนพฤษภาคม 2557 จนถึงเดือนธันวาคม 2559
และเป็นผู้นำรัฐบาลพลเรือนตั้งแต่เดือนธันวาคม 2559 จนถึงปัจจุบัน เขาเป็นผู้นำรัฐบาลที่ดำรงตำแหน่งยาวนานที่สุดในประวัติศาสตร์ของไทย
"""

Training hyperparameters

The following hyperparameters were used during training:

learning_rate: 0.0002
train_batch_size: 8
eval_batch_size: 8
seed: 42
optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
lr_scheduler_type: cosine
mixed_precision_training: Native AMP

Framework versions

PEFT 0.7.1
Transformers 4.37.0.dev0
Pytorch 2.1.2+cu121
Datasets 2.16.0
Tokenizers 0.15.0