tosin2013 commited on
Commit
eb872b8
·
1 Parent(s): 7e1e3c5

adding zerogpu

Browse files
Files changed (1) hide show
  1. app.py +9 -4
app.py CHANGED
@@ -37,18 +37,23 @@ if model_provider.lower() == "openai":
37
  )
38
  else:
39
  MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"
40
- # Initialize Hugging Face InferenceClient
41
  hf_client = InferenceClient(
42
  model=MODEL_NAME,
43
- api_key=os.environ.get("HF_TOKEN")
 
44
  )
45
 
46
  # Load the Hugging Face dataset
47
  dataset = load_dataset('tosin2013/autogen', streaming=True)
48
  dataset = Dataset.from_list(list(dataset['train']))
49
 
50
- # Initialize embeddings
51
- embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
 
 
 
52
 
53
  # Extract texts from the dataset
54
  texts = dataset['input']
 
37
  )
38
  else:
39
  MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"
40
+ # Initialize Hugging Face InferenceClient with GPU support
41
  hf_client = InferenceClient(
42
  model=MODEL_NAME,
43
+ api_key=os.environ.get("HF_TOKEN"),
44
+ timeout=120 # Increased timeout for GPU inference
45
  )
46
 
47
  # Load the Hugging Face dataset
48
  dataset = load_dataset('tosin2013/autogen', streaming=True)
49
  dataset = Dataset.from_list(list(dataset['train']))
50
 
51
+ # Initialize embeddings with GPU support if available
52
+ device = "cuda" if torch.cuda.is_available() else "cpu"
53
+ embeddings = HuggingFaceEmbeddings(
54
+ model_name="sentence-transformers/all-MiniLM-L6-v2",
55
+ model_kwargs={"device": device}
56
+ )
57
 
58
  # Extract texts from the dataset
59
  texts = dataset['input']