ml_exercise / llm_classification.py
trantuan1701's picture
hehe
ba48a36
raw
history blame contribute delete
914 Bytes
from llm import llm
from langchain.prompts import PromptTemplate
from langchain_core.messages import SystemMessage, HumanMessage
def get_answer(q: str) -> int:
system_prompt = (
"You are an assistant classifying a sentence sentiment. "
"Output strictly '1' for positive or '0' for negative. No extra text."
)
user_prompt = PromptTemplate(
input_variables=["q"],
template=(
"Classify this sentence: {q}\n"
"Output only 1 if positive or 0 if negative."
),
)
try:
resp = llm.invoke([
SystemMessage(content=system_prompt),
HumanMessage(content=user_prompt.format(q=q))
])
text = resp.content.strip() # lấy nội dung model trả về
return 1 if text == "1" else 0
except Exception as e:
print(f"Error: {e}")
return 0 # hoặc None tùy bạn muốn