Views
No views yet
1from transformers import GPT2LMHeadModel, GPT2Tokenizer
2
3# Load tokenizer and model
4tokenizer = GPT2Tokenizer.from_pretrained("RaushanTurganbay/GPT2_instruct_tuned")
5model = GPT2LMHeadModel.from_pretrained("RaushanTurganbay/GPT2_instruct_tuned")
6
7# Generate responses
8class StoppingCriteriaSub(StoppingCriteria):
9 def __init__(self, stops=[], encounters=1):
10 super().__init__()
11 self.stops = [stop.to("cuda") for stop in stops]
12 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
13 for stop in self.stops:
14 if torch.all((stop == input_ids[0][-len(stop):])).item():
15 return True
16 return False
17
18
19def stopping_criteria(tokenizer, stop_words):
20 stop_words_ids = [tokenizer(stop_word, return_tensors='pt')['input_ids'].squeeze() for stop_word in stop_words]
21 stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])
22 return stopping_criteria
23
24# Generate responses
25stopping = stopping_criteria(tokenizer, ["\n\nHuman:"])
26prompt = "\n\nHuman: {your_instruction}\n\nAssistant:"
27inputs = tokenizer(prompt, return_tensors="pt")
28outputs = model.generate(**inputs, stopping_criteria=stopping, max_length=150)
29
30print("Model Response:", tokenizer.batch_decode(outputs))