The Rise of Domestic AI: Integrating Zhipu AI with LlamaIndex for Data Processing

The rise of domestic AI is indeed a long and arduous journey.

As mentioned earlier, Llama Index is a simple flexible data framework. Since the official integration is not available yet, let’s integrate it ourselves.

Integrating Zhipu AI

LlamaIndex provides an abstract class called CustomLLM to facilitate the integration of our own large model products:

class CustomLLM(LLM):
    """Simple abstract base class for custom LLMs.

    Subclasses must implement the `__init__`, `complete`,
        `stream_complete`, and `metadata` methods.
    """

From the comments, we can see that implementing our own LLM is actually quite simple; we just need to implement the __init__, complete, stream_complete, and metadata methods.

Metadata Method

What is Metadata?Metadata is formally defined as data that describes other data. It can be a bit convoluted, so it requires careful consideration. In simple terms, it refers to the attributes of records, the structure of tables, the layout of interfaces, the characteristics of data, and the parameters of programs.

Among the four methods, the one we might be less familiar with is metadata. The metadata of an LLM generally includes the following:

• • model_name: str = "chatglm_turbo"

• • top_p = 0.7

• • temperature = 0.9

• • num_output: int = 256

• • context_window: int = 3900

Parameters like “model name,” “temperature,” and “top p” are likely already familiar to everyone. Let’s implement it simply:

class ZhipuLLM(CustomLLM):
    model_name: str = "chatglm_turbo"
    top_p = 0.7
    temperature = 0.9

    @property
    def metadata(self) -> LLMMetadata:
        """Get LLM metadata."""
        return LLMMetadata(
            context_window=self.context_window,
            num_output=self.num_output,
            model_name=self.model_name,
        )

Since Zhipu AI’s API doesn’t have too many parameters, we can directly use reasonable values for top_p and temperature in the calling program. In fact, we could even implement an empty method for metadata:

class ZhipuLLM(CustomLLM):
    @property
    def metadata(self) -> LLMMetadata:
        return LLMMetadata()

Complete Method

Next, we implement the complete method. Previously, we had an invoke_prompt method to implement the call to Zhipu AI. Here, we can use it directly without instantiating or setting too many parameters; the simpler, the better:

@llm_completion_callback()
def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
    response = invoke_prompt(prompt)
    return CompletionResponse(text=response)

stream_completeStreamed returns have examples from the official documentation, so I won’t implement that here.

Text Embedding

Zhipu AI also provides a text embedding interface, so we will also use it directly, thus completely freeing ourselves from OpenAI and the need to build our own model.

The embedding in LlamaIndex is based on BaseEmbedding, so we will also write an implementation for Zhipu AI. There isn’t enough time for extensive explanations, so let’s look at the code directly:

class ZhipuEmbedding(BaseEmbedding):
    _model: str = PrivateAttr()
    _instruction: str = PrivateAttr()

    def __init__(
            self,
            instructor_model_name: str = "text_embedding",
            instruction: str = "Represent a document for semantic search:",
            **kwargs: Any,
    ) -> None:
        self._model = 'text_embedding'
        self._instruction = instruction
        super().__init__(**kwargs)

    @classmethod
    def class_name(cls) -> str:
        return "zhipu_embeddings"

    async def _aget_query_embedding(self, query: str) -> List[float]:
        return self._get_query_embedding(query)

    async def _aget_text_embedding(self, text: str) -> List[float]:
        return self._get_text_embedding(text)

    def _get_query_embedding(self, query: str) -> List[float]:
        embeddings = invoke_embedding(query)
        return embeddings

    def _get_text_embedding(self, text: str) -> List[float]:
        embeddings = invoke_embedding(text)
        return embeddings

    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        return [self._get_text_embedding(text) for text in texts]

LlamaIndex Combined with Zhipu AI

Now let’s put together the results from earlier and write a simple program to see if it works:

First, we need to place a text file in the directory. PDF, Excel, Docx files are all acceptable. Here, I will use the previous design document for the Asian Games medals:

Then, we can directly run the program:

# define our LLM
llm = ZhipuLLM()
embed_model = ZhipuEmbedding()

service_context = ServiceContext.from_defaults(
    llm=llm, embed_model=embed_model
)

# Load your data
documents = SimpleDirectoryReader(doc_path).load_data()
index = VectorStoreIndex.from_documents(documents, service_context=service_context)

# Query and print response
query_engine = index.as_query_engine()
response = query_engine.query("What are the processes and design concepts for the medals of the Hangzhou Asian Games? Please answer in Chinese.")
print(response)

Let’s look at the answer:

The processes and design concepts for the medals of the Hangzhou Asian Games are as follows:

1. Process: The medals use a series of complex processes, including stamping, milling shapes and grooves, finishing and polishing, gold and silver plating, and surface protection. The ribbon adopts jacquard weaving technology and environmentally friendly printing technology, with double-sided hand stitching. Additionally, the front of the medal features a traditional Jiangnan craft—copper engraving, combining the four calligraphic styles (regular, cursive, clerical, and seal) with the English abbreviation of the Hangzhou Asian Games.

2. Design Concept: The medal design integrates Hangzhou's three world cultural heritage sites—West Lake, Grand Canal, and Liangzhu Ancient City Ruins, showcasing Hangzhou's landscape and Jiangnan culture. The front of the medal outlines the painting of Hangzhou city, while the back resembles a square seal, symbolizing the beautiful mark left by athletes at the Hangzhou Asian Games. The overall design is unique and highly recognizable, embodying the meaning of harmony and diversity.

In addition, the medal design also reflects Hangzhou's character as a city of ecological civilization, with beautiful lakes and mountains, green waters and hills, creating a golden mountain and silver mountain, while embodying the spirit of striving for excellence in sports. The medal has a diameter of 70mm and a thickness of 6mm, with a gradient color on the front and a seal pattern presented in Hangzhou Jiangnan embroidery technology on the back. The overall design of the medal is full of Eastern aesthetic beauty, combining the rich heritage of Southern Song culture with the regional characteristics of Hangzhou, expressing a unique aesthetic philosophy.

The answers are all sourced from the document but have been reorganized. I just want to say that the expression is so fluent in Chinese! Far surpassing my own.

Conclusion

The use of large models in Llama Index mainly occurs in the two areas mentioned above. Once these two points are connected, the remaining various applications will be no different from using OpenAI. I will share more scenarios and practical AI applications later.

References

[1] Zhipu AI: https://www.zhipuai.cn/[2] Open Platform: https://open.bigmodel.cn/[3] API Usage Guide: https://open.bigmodel.cn/dev/api[4] LlamaIndex: https://www.llamaindex.ai/[5] LlamaIndex Documentation: https://docs.llamaindex.ai/en/stable/[6] Available LLM integrations: https://docs.llamaindex.ai/en/stable/module_guides/models/llms/modules.html

— END —