poem2pic / t5_summary.py
barunsaha's picture
Add application files
854eca0
raw
history blame
939 Bytes
from langchain import HuggingFaceHub, PromptTemplate
from config import Config
llm = HuggingFaceHub(
repo_id=Config.LLM_MODEL_NAME,
huggingfacehub_api_token=Config.HUGGINGFACEHUB_API_TOKEN,
model_kwargs={'temperature': Config.LLM_TEMPERATURE, 'max_length': Config.LLM_MAX_OUTPUT_LENGTH}
)
# print(llm)
template = '''
You are a helpful assistant who can answer questions about literature and poetry. You are creative.
Consider the following poem:
{poem}
Write a creative summary of the poem in about 20 words.
The summary should try to capture the subject, objects, key themes, and sentiments.
'''
prompt = PromptTemplate.from_template(template)
def generate_summary(poem: str) -> str:
"""
Generate the summary of a poem using an LLM.
:param poem: The text of the poem
:return: The summary
"""
print(prompt.format(poem=poem))
summary = llm(prompt.format(poem=poem))
return summary