Skip to content

Instantly share code, notes, and snippets.

@gante

gante/galactica_contrastive_search.py

Created November 18, 2022 11:33
Show Gist options
  • Save gante/fc67e189463178f2716d0be7df463661 to your computer and use it in GitHub Desktop.
Save gante/fc67e189463178f2716d0be7df463661 to your computer and use it in GitHub Desktop.
Download ZIP
Galactica (1.3b) + contrastive search examples
Raw
galactica_contrastive_search.py
from transformers import AutoTokenizer, OPTForCausalLM
tokenizer = AutoTokenizer.from_pretrained("facebook/galactica-1.3b")
model = OPTForCausalLM.from_pretrained("facebook/galactica-1.3b", device_map="auto")
# input_text = "Question: How small is a human cell? Answer:" # they should get the same short answers
input_text = "Question: What do Maxwell's equations represent? Answer:" # better with repetitions
# input_text = "Question: Simplify the following Python code using math:```pythondef calc_sum(n): i = 0 s = 0 while i <= n: s += i i += 1 return s```Answer:" # better with early stop
# input_text = "Question: What technology will revolutionize language models? Answer:"
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
outputs = model.generate(input_ids, max_new_tokens=128)
print("Greedy:\n", tokenizer.decode(outputs[0]))
outputs = model.generate(input_ids, max_new_tokens=128, top_k=4, penalty_alpha=0.6)
print("\nContrastive:\n", tokenizer.decode(outputs[0]))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment