Update app.py

app.py

@@ -2,6 +2,10 @@ import torch
 
 import gradio as gr
 import argparse
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+# from transformers.generation.utils import GenerationConfig
+
 from utils import load_hyperparam, load_model
 from models.tokenize import Tokenizer
 from models.llama import *
@@ -36,22 +40,25 @@ def init_args():
 
     args = load_hyperparam(args)
 
-    args.tokenizer = Tokenizer(model_path=args.spm_model_path)
+    # args.tokenizer = Tokenizer(model_path=args.spm_model_path)
+    args.tokenizer = AutoTokenizer.from_pretrained("Linly-AI/Linly-ChatFlow", use_fast=False, trust_remote_code=True)
     args.vocab_size = args.tokenizer.sp_model.vocab_size()
 
 
 def init_model():
     global lm_generation
-    torch.set_default_tensor_type(torch.HalfTensor)
-    model = LLaMa(args)
-    torch.set_default_tensor_type(torch.FloatTensor)
+    # torch.set_default_tensor_type(torch.HalfTensor)
+    # model = LLaMa(args)
+    # torch.set_default_tensor_type(torch.FloatTensor)
+    # # args.load_model_path = hf_hub_download(repo_id=args.load_model_path, filename='chatflow_13b.bin')
     # args.load_model_path = hf_hub_download(repo_id=args.load_model_path, filename='chatflow_13b.bin')
-    args.load_model_path = hf_hub_download(repo_id=args.load_model_path, filename='chatflow_13b.bin')
-    model = load_model(model, args.load_model_path)
-    model.eval()
+    # model = load_model(model, args.load_model_path)
+    # model.eval()
 
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
+    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # model.to(device)
+    model = AutoModelForCausalLM.from_pretrained("Linly-AI/Linly-ChatFlow", device_map="auto", torch_dtype=torch.float16, trust_remote_code=True)
+    
     print(torch.cuda.max_memory_allocated() / 1024 ** 3)
     lm_generation = LmGeneration(model, args.tokenizer)