From f1afae7b26ceaa040adcd4b4812b64f259b05647 Mon Sep 17 00:00:00 2001
From: Stanislas0 <qinkai.zheng1028@gmail.com>
Date: Tue, 7 Mar 2023 21:46:53 +0800
Subject: [PATCH] Remove dependency

---
 codegeex/__init__.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/codegeex/__init__.py b/codegeex/__init__.py
index 954e66b..0bc55c6 100644
--- a/codegeex/__init__.py
+++ b/codegeex/__init__.py
@@ -1,7 +1,6 @@
 import copy
 
 from typing import *
-from codegeex.megatron.model import CodeGeeXModel
 from codegeex.tokenizer import CodeGeeXTokenizer
 from codegeex.torch.inference import get_token_stream
 
@@ -14,7 +13,7 @@ def get_model(
 
 
 def generate(
-    model: CodeGeeXModel, 
+    model, 
     tokenizer: CodeGeeXTokenizer, 
     prompt: str, 
     out_seq_length: int,