update

app.py

@@ -2,7 +2,7 @@ import os
 import gradio as gr
 
 import torch
-from transformers import pipeline, StoppingCriteriaList, MaxTimeCriteria
+from transformers import pipeline, StoppingCriteriaList, MaxTimeCriteria, AutoTokenizer, AutoModelForCausalLM
 from openai import OpenAI
 
 from elasticsearch import Elasticsearch
@@ -79,8 +79,11 @@ def rag_pipeline(prompt, index="pubmed", num_docs=3, model_name="HuggingFaceH4/z
         print('OAI_RESPONSE', openai_res)
         response = openai_res.choices[0].message.content.strip()
     else:
+        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+
         # Load your language model from HuggingFace Transformers
-        generator = pipeline("text-generation", model=model_name, torch_dtype=torch.bfloat16, device_map="auto")
+        generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
         # Generate response using the HF LLM
         hf_response = generator(messages, **generation_kwargs)