init: first memary/synonym_expand

llama-index-packs/llama-index-packs-memary/llama_index/packs/memary/synonym_expand/__init__.py

@@ -0,0 +1,4 @@
+from llama_index.packs.memary.synonym_expand.synonym import custom_synonym_expand_fn
+from llama_index.packs.memary.synonym_expand.output import SynonymOutput
+
+__all__ = ["custom_synonym_expand_fn", "SynonymOutput"]

llama-index-packs/llama-index-packs-memary/llama_index/packs/memary/synonym_expand/output.py

@@ -0,0 +1,10 @@
+from typing import List
+
+from langchain_core.pydantic_v1 import BaseModel, Field
+
+
+class SynonymOutput(BaseModel):
+ synonyms: List[str] = Field(
+ description=
+ "Synonyms of keywords provided, make every letter lowercase except for the first letter"
+ )

llama-index-packs/llama-index-packs-memary/llama_index/packs/memary/synonym_expand/synonym.py

@@ -0,0 +1,49 @@
+from langchain_openai import OpenAI
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import JsonOutputParser
+from typing import List
+import os
+from llama_index.packs.memary.synonym_expand.output import SynonymOutput
+from dotenv import load_dotenv
+
+
+def custom_synonym_expand_fn(keywords: str) -> List[str]:
+ load_dotenv()
+ llm = OpenAI(openai_api_key=os.getenv("OPENAI_KEY"), temperature=0)
+ parser = JsonOutputParser(pydantic_object=SynonymOutput)
+
+ template = """
+ You are an expert synonym exapnding system. Find synonyms or words commonly used in place to reference the same word for every word in the list:
+
+ Some examples are:
+ - a synonym for Palantir may be Palantir technologies or Palantir technologies inc.
+ - a synonym for Austin may be Austin texas
+ - a synonym for Taylor swift may be Taylor 
+ - a synonym for Winter park may be Winter park resort
+
+ Format: {format_instructions}
+
+ Text: {keywords}
+ """
+
+ prompt = PromptTemplate(
+ template=template,
+ input_variables=["keywords"],
+ partial_variables={
+ "format_instructions": parser.get_format_instructions()
+ },
+ )
+
+ chain = prompt | llm | parser
+ result = chain.invoke({"keywords": keywords})
+
+ l = []
+ for category in result:
+ for synonym in result[category]:
+ l.append(synonym.capitalize())
+
+ return l
+
+
+# testing
+# print(custom_synonym_expand_fn("[Nvidia]"))