Alberto Carmona commited on
Commit
47d82ab
·
1 Parent(s): f90a30f

Setn env var for cuda

Browse files
Files changed (1) hide show
  1. functions.py +12 -7
functions.py CHANGED
@@ -1,9 +1,13 @@
 
 
1
  import requests
2
- from bs4 import BeautifulSoup
3
  import torch
4
- from peft import PeftModel, PeftConfig
 
5
  from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
6
 
 
 
7
 
8
  generation_config = GenerationConfig(temperature=.8,
9
  top_p=0.75,
@@ -27,10 +31,10 @@ def summarize_text(text: str):
27
  batch = tokenizer(input_text, return_tensors='pt')
28
  print(['summarize_text', 'generating'])
29
  with torch.cuda.amp.autocast():
30
- output_tokens = model.generate(**batch,
31
- max_new_tokens=256,
32
- generation_config=generation_config
33
- )
34
  output = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
35
  print(['summarize_text', 'end'])
36
  return output
@@ -50,4 +54,5 @@ def load_model(peft_model_id):
50
  return model, tokenizer
51
 
52
 
53
- model, tokenizer = load_model("hackathon-somos-nlp-2023/opt-6.7b-lora-sag-t3000-v300-v2")
 
 
1
+ import os
2
+
3
  import requests
 
4
  import torch
5
+ from bs4 import BeautifulSoup
6
+ from peft import PeftConfig, PeftModel
7
  from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
8
 
9
+ os.environ["CUDA_VISIBLE_DEVICES"] = "0"
10
+
11
 
12
  generation_config = GenerationConfig(temperature=.8,
13
  top_p=0.75,
 
31
  batch = tokenizer(input_text, return_tensors='pt')
32
  print(['summarize_text', 'generating'])
33
  with torch.cuda.amp.autocast():
34
+ output_tokens = model.generate(**batch,
35
+ max_new_tokens=256,
36
+ generation_config=generation_config
37
+ )
38
  output = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
39
  print(['summarize_text', 'end'])
40
  return output
 
54
  return model, tokenizer
55
 
56
 
57
+ model, tokenizer = load_model(
58
+ "hackathon-somos-nlp-2023/opt-6.7b-lora-sag-t3000-v300-v2")