rasyosef commited on
Commit
1d57a43
1 Parent(s): 0858c92

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -37,7 +37,7 @@ QA_PROMPT = PromptTemplate(
37
  model_id = "microsoft/phi-2"
38
 
39
  tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
40
- model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, device_map="auto", trust_remote_code=True)
41
 
42
  # sentence transformers to be used in vector store
43
  embeddings = HuggingFaceEmbeddings(
@@ -86,7 +86,7 @@ def generate(question, answer, text_file, max_new_tokens):
86
  phi2_pipeline = pipeline(
87
  "text-generation", tokenizer=tokenizer, model=model, max_new_tokens=max_new_tokens,
88
  pad_token_id=tokenizer.eos_token_id, eos_token_id=tokenizer.eos_token_id,
89
- device_map="auto", streamer=streamer
90
  )
91
 
92
  hf_model = HuggingFacePipeline(pipeline=phi2_pipeline)
 
37
  model_id = "microsoft/phi-2"
38
 
39
  tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
40
+ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, device_map="cpu", trust_remote_code=True)
41
 
42
  # sentence transformers to be used in vector store
43
  embeddings = HuggingFaceEmbeddings(
 
86
  phi2_pipeline = pipeline(
87
  "text-generation", tokenizer=tokenizer, model=model, max_new_tokens=max_new_tokens,
88
  pad_token_id=tokenizer.eos_token_id, eos_token_id=tokenizer.eos_token_id,
89
+ device_map="cpu", streamer=streamer
90
  )
91
 
92
  hf_model = HuggingFacePipeline(pipeline=phi2_pipeline)