Update app.py

app.py

@@ -1,14 +1,15 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
 from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 from threading import Thread
 
 # Loading the tokenizer and model from Hugging Face's model hub.
-# model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-model_name = "scutcyr/BianQue-2"
-tokenizer = AutoTokenizer.from_pretrained(model_name,trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name,trust_remote_code=True)
+# model_name_or_path = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+model_name_or_path = "scutcyr/BianQue-2"
+tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,trust_remote_code=True)
+# model = AutoModelForCausalLM.from_pretrained(model_name,trust_remote_code=True)
+model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True).half()
 
 # using CUDA for an optimal experience
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')