hanzla commited on
Commit
f38df4a
·
1 Parent(s): 851843d
Files changed (1) hide show
  1. src/pdfchatbot.py +3 -3
src/pdfchatbot.py CHANGED
@@ -58,8 +58,8 @@ class PDFChatBot:
58
  self.pipeline = pipeline(
59
  "text-generation",
60
  model="meta-llama/Meta-Llama-3-8B-Instruct",
61
- model_kwargs={"torch_dtype": torch.float16},
62
- device_map="auto",
63
  )
64
  print("Model pipeline loaded")
65
 
@@ -69,7 +69,7 @@ class PDFChatBot:
69
  self.current_context = context
70
  print("Context Ready")
71
  print(self.current_context)
72
-
73
  def create_organic_response(self, history, query):
74
  self.get_organic_context(query)
75
  messages = [
 
58
  self.pipeline = pipeline(
59
  "text-generation",
60
  model="meta-llama/Meta-Llama-3-8B-Instruct",
61
+ model_kwargs={"torch_dtype": torch.bfloat16},
62
+ device="cuda",
63
  )
64
  print("Model pipeline loaded")
65
 
 
69
  self.current_context = context
70
  print("Context Ready")
71
  print(self.current_context)
72
+ @spaces.GPU
73
  def create_organic_response(self, history, query):
74
  self.get_organic_context(query)
75
  messages = [