from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_community.llms import CTransformers
import os
# Cài thư việt trước
# ! pip install langchain langchain-community ctransformers

MODEL_FILE = "Cube-Python.gguf"
MODEL_TYPE = "llama"
GPU_LAYERS = 0 # Nếu máy có GPU VRAM nhiều hơn có thể chỉnh lên 10-20 để trả lời nhanh hơn
CONTEXT_LENGTH = 4096


def load_llm_with_ctransformers():
    model_path = os.path.join(os.getcwd(), MODEL_FILE)

    if not os.path.exists(model_path):
        raise FileNotFoundError(f"File mô hình {MODEL_FILE} không tồn tại.")

    llm = CTransformers(
        model=model_path,
        model_type=MODEL_TYPE,
        config={
            'max_new_tokens': 1024,
            'temperature': 0.1,
            'gpu_layers': GPU_LAYERS,
            'context_length': CONTEXT_LENGTH,
        }
    )
    return llm


template = """[INST] Bạn là một trợ lý AI chuyên nghiệp về lập trình Python.
Hãy viết code Python chất lượng cao để giải quyết yêu cầu sau.
Ưu tiên:
- code rõ ràng
- có thể chạy được
- đặt tên biến dễ hiểu
Chỉ trả lời bằng code.
Yêu cầu: {question} [/INST]"""

prompt = PromptTemplate(
    input_variables=["question"],
    template=template
)

llm = load_llm_with_ctransformers()

parser = StrOutputParser()

chain = prompt | llm | parser

question = '''
Write a Python program that takes a sentence as input and removes all punctuation marks
from it.
Input:
A sentence: "Hello! This is a NLP practical exam, isn't it?"
Desired Output:
Hello This is a NLP practical exam isnt it'''

result = chain.invoke({"question": question})

print(result)