Update app.py

app.py

@@ -2,14 +2,21 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
 from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
+from transformers.generation.utils import GenerationConfig
 from threading import Thread
 
 # Loading the tokenizer and model from Hugging Face's model hub.
 # model_name_or_path = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 model_name_or_path = "Flmc/DISC-MedLLM"
-tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,trust_remote_code=True)
+# 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)
+# model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True)
+
+
+tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False, trust_remote_code=True)
+ model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True)
+model.generation_config = GenerationConfig.from_pretrained(model_name_or_path)
+
 
 # using CUDA for an optimal experience
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')