alx-d commited on
Commit
a5abbb0
·
verified ·
1 Parent(s): 24cbe1e

Update philosophy.py

Browse files
Files changed (1) hide show
  1. philosophy.py +4 -3
philosophy.py CHANGED
@@ -1,6 +1,7 @@
 
1
  from transformers import AutoModelForCausalLM, AutoTokenizer
 
2
  from llama_index import (
3
- LLMPredictor,
4
  PromptHelper,
5
  StorageContext,
6
  ServiceContext,
@@ -33,7 +34,7 @@ def create_llm_pipeline():
33
  "text-generation",
34
  model=model,
35
  tokenizer=tokenizer,
36
- device=DEVICE,
37
  max_length=2048,
38
  do_sample=True,
39
  temperature=0.7,
@@ -121,4 +122,4 @@ if __name__ == "__main__":
121
  # Create initial index
122
  index = data_ingestion_indexing("books")
123
  # Launch the interface
124
- iface.launch()
 
1
+ import transformers # Added since we use transformers.pipeline below
2
  from transformers import AutoModelForCausalLM, AutoTokenizer
3
+ from llama_index.llm_predictor import LLMPredictor # Updated import path for LLMPredictor
4
  from llama_index import (
 
5
  PromptHelper,
6
  StorageContext,
7
  ServiceContext,
 
34
  "text-generation",
35
  model=model,
36
  tokenizer=tokenizer,
37
+ device=0 if DEVICE == "cuda" else -1, # Use device index: 0 for GPU, -1 for CPU
38
  max_length=2048,
39
  do_sample=True,
40
  temperature=0.7,
 
122
  # Create initial index
123
  index = data_ingestion_indexing("books")
124
  # Launch the interface
125
+ iface.launch()