tasksource-deberta-small-long-nli / llm_classifier_service.py
shomez's picture
Upload llm_classifier_service.py with huggingface_hub
11716e5 verified
from langchain_community.llms.ollama import Ollama
# from langchain_community.llms import Ollama
from langchain_core.prompts import PromptTemplate
PROMPT_TEMPLATE = """Examples
Document: USC is one of the top universities in America
Topic: Universities
Answer: Yes
Document: USC is one of the top universities in America
Topic: Politics
Answer: No
Document: Governments worldwide are grappling with the challenges of regulating artificial intelligence to ensure it benefits society while protecting individual privacy.
Topic: Politics
Answer: Yes
Answer if the Topic Classification of the following Document is correct.
Answer only in "Yes" or "No".
Document: {document}
Topic: {topic}
Answer:"""
def topic_in_text_service(document: str, topic: str, model_id: str, base_url: str) -> bool:
llm = Ollama(model="llama3", base_url="http://ollama:11434")
prompt_template = PromptTemplate.from_template(PROMPT_TEMPLATE)
prompt = prompt_template.format(document=document, topic=topic)
return 'yes' in llm.invoke(prompt, num_predict=1).lower()