Update src/models/chatglm2.py

src/models/chatglm2.py

@@ -1,24 +0,0 @@
-# src/models/chatglm2.py
-from transformers import AutoTokenizer, AutoModel
-import torch
-import os
-
-class ChatGLM2:
-    def __init__(self):
-        model_id = "THUDM/chatglm2-6b-int4"
-        offload_path = "./offload_chatglm2"
-        os.makedirs(offload_path, exist_ok=True)
-
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-        self.model = AutoModel.from_pretrained(
-            model_id,
-            trust_remote_code=True,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            offload_folder=offload_path
-        ).eval()
-
-    def generate(self, prompt):
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
-        outputs = self.model.generate(**inputs, max_new_tokens=256)
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)