Spaces:
Sleeping
Sleeping
File size: 1,067 Bytes
8f53255 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
from langchain_nvidia_ai_endpoints import ChatNVIDIA
from langchain_core.prompts import PromptTemplate
from dotenv import load_dotenv
import os
load_dotenv()
template = """
You are a helpful Summarizer Chatbot who can just summarize the input text and return.
User : {input_text}
AI : """
prompt = PromptTemplate.from_template(template)
class Summarizer(object):
def __init__(self) -> None:
self.llm = ChatNVIDIA(
model='google/gemma-2-2b-it',
api_key=os.getenv('NV_API_KEY'),
max_tokens=128,
temperature=0.01,
top_p=.7
)
self.chain = prompt | self.llm
def summarize(self, mem:list):
summarized_memory = []
for item in mem:
if item['role'].lower() == 'user':
summarized_memory.append(item)
else:
summarized_content = self.chain.invoke(item['content']).content
summarized_memory.append({'role' : item['role'], 'content' : summarized_content})
return summarized_memory
|