Tijmen2 commited on
Commit
fbadf1a
1 Parent(s): bb103b3

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +14 -16
README.md CHANGED
@@ -30,24 +30,22 @@ The code used to generate cosmosage_v2 is available at https://github.com/tijmen
30
  After downloading cosmosage_v2, the following example code can be used to ask questions:
31
 
32
  ```python
33
- path_to_model = 'cosmosage_v2/'
34
 
35
- from transformers import AutoModelForCausalLM, AutoTokenizer
36
  import torch
37
- device = "cuda"
38
- model = AutoModelForCausalLM.from_pretrained(path_to_model).to(device)
39
- tokenizer = AutoTokenizer.from_pretrained(path_to_model)
40
- def ask_cosmosage(question):
41
- input_ids = torch.cat([
42
- tokenizer.encode("You are cosmosage, an AI programmed to be a cosmology expert. You answer the USER's question clearly in long form, always providing context. When appropriate, provide a reference.", return_tensors="pt"),
43
- torch.tensor([[28705]]),
44
- tokenizer.encode("USER:", add_special_tokens=False, return_tensors="pt"),
45
- tokenizer.encode(question, add_special_tokens=False, return_tensors="pt"),
46
- torch.tensor([[28705]]),
47
- tokenizer.encode("ASSISTANT:", add_special_tokens=False, return_tensors="pt")
48
- ], dim=-1).to(device)
49
- generated_ids = model.generate(input_ids, max_length=input_ids.shape[1] + 1000, do_sample=True, temperature=0.4)
50
- return tokenizer.decode(generated_ids[0], skip_special_tokens=True)
51
  ```
52
 
53
  ## Comparison to cosmosage_v1
 
30
  After downloading cosmosage_v2, the following example code can be used to ask questions:
31
 
32
  ```python
33
+ model_path = "models/cosmosage_v2/"
34
 
 
35
  import torch
36
+ from transformers import AutoModelForCausalLM, AutoTokenizer
37
+
38
+ device = torch.device("cuda")
39
+ model = AutoModelForCausalLM.from_pretrained(model_path).to(device, dtype=torch.bfloat16)
40
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
41
+
42
+ def ask_cosmosage(question, answer_start=''):
43
+ prompt = f"You are cosmosage, an AI programmed to be a cosmology expert. You answer the USER's question clearly in long form, always providing context. When appropriate, provide a reference.USER: {question}ASSISTANT: {answer_start}"
44
+ input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
45
+ generated_ids = model.generate(input_ids, max_length=1024, do_sample=True, temperature=0.7, top_k=None, pad_token_id=tokenizer.eos_token_id)
46
+ generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
47
+ answer = generated_text.split("ASSISTANT:")[-1]
48
+ return answer
 
49
  ```
50
 
51
  ## Comparison to cosmosage_v1