Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Changes to make this compatible with running purely on a GPU #56

Closed
wants to merge 5 commits into from
+21 −15
Closed

Changes to make this compatible with running purely on a GPU #56

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
27b5eb4
Changes to make this compatible with running purely on a GPU
rishsriv Dec 8, 2023
45bc3df
Update hf_runner.py
rishsriv Dec 8, 2023
1a82bf0
Update hf_runner.py
rishsriv Dec 8, 2023
0db685d
initial commit
wendy-aw Dec 11, 2023
4d86d7f
fix device of column_emb
wendy-aw Dec 11, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 4 additions & 13 deletions eval/hf_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,21 +61,12 @@ def get_tokenizer_model(model_name: Optional[str], adapter_path: Optional[str]):
model = PeftModel.from_pretrained(model, adapter_path)
model = model.merge_and_unload()
print(f"Merged adapter {adapter_path}")
elif model_name is not None and "llama" in model_name:
print(f"Loading Llama-based model {model_name}")
tokenizer = LlamaTokenizer.from_pretrained(
model_name, legacy=False, use_fast=True
)
model = LlamaForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="auto",
use_cache=True,
use_flash_attention_2=True,
)
else:
print(f"Loading model {model_name}")
tokenizer = AutoTokenizer.from_pretrained(model_name)
try:
tokenizer = AutoTokenizer.from_pretrained(model_name)
except:
tokenizer = AutoTokenizer.from_pretrained("codellama/CodeLlama-34b-hf")
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
Expand Down
11 changes: 11 additions & 0 deletions prompts/prompt2.md
View file
Open in desktop
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Checking if this an improved version of prompt.md? If so, shall we just replace the existing prompt.md with this?

Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
### Instructions
Generate a SQL query to answer the following question:
`{user_question}`

### Schema
The query will run on a database with the following schema:
{table_metadata_string}

### SQL
Here is a query that answers the question `{user_question}`
```
8 changes: 6 additions & 2 deletions utils/pruning.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
if os.getenv("TOKENIZERS_PARALLELISM") is None:
os.environ["TOKENIZERS_PARALLELISM"] = "false"

encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device="cpu")
encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was thinking of having a variable at the top of pruning.py called user_device, and then passing it to everywhere torch device is required. We can just check and set it once instead of checking/setting every place we create a tensor.

# at the top
user_device = "cuda" if torch.cuda.is_available() else "cpu"
encoder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device=user_device)
...
query_emb = encoder.encode(query, convert_to_tensor=True, device=user_device)
...
column_emb = column_emb.to(user_device)

The other benefit is that we can easily override user_device in tests to run only on cpu. wdyt?

nlp = spacy.load("en_core_web_sm")


Expand All @@ -23,7 +23,7 @@ def knn(
"""
Get top most similar columns' embeddings to query using cosine similarity.
"""
query_emb = encoder.encode(query, convert_to_tensor=True, device="cpu").unsqueeze(0)
query_emb = encoder.encode(query, convert_to_tensor=True).unsqueeze(0)
similarity_scores = F.cosine_similarity(query_emb, all_emb)
top_results = torch.nonzero(similarity_scores > threshold).squeeze()
# if top_results is empty, return empty tensors
Expand Down Expand Up @@ -95,6 +95,10 @@ def get_md_emb(
3. Generate the metadata string using the column info so far.
4. Get joinable columns between tables in topk_table_columns and add to final metadata string.
"""
if torch.cuda.is_available():
column_emb = column_emb.to("cuda")
else:
column_emb = column_emb.to("cpu")
# 1) get top k columns
top_k_scores, top_k_indices = knn(question, column_emb, k, threshold)
topk_table_columns = {}
Expand Down