Qwen2.5-Coder-repolevel.py
1 from transformers import AutoTokenizer, AutoModelForCausalLM 2 device = "cuda" # the device to load the model onto 3 4 # Now you do not need to add "trust_remote_code=True" 5 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-32B") 6 model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-Coder-32B", device_map="auto").eval() 7 8 # tokenize the input into tokens 9 input_text = """<|repo_name|>library-system 10 <|file_sep|>library.py 11 class Book: 12 def __init__(self, title, author, isbn, copies): 13 self.title = title 14 self.author = author 15 self.isbn = isbn 16 self.copies = copies 17 18 def __str__(self): 19 return f"Title: {self.title}, Author: {self.author}, ISBN: {self.isbn}, Copies: {self.copies}" 20 21 class Library: 22 def __init__(self): 23 self.books = [] 24 25 def add_book(self, title, author, isbn, copies): 26 book = Book(title, author, isbn, copies) 27 self.books.append(book) 28 29 def find_book(self, isbn): 30 for book in self.books: 31 if book.isbn == isbn: 32 return book 33 return None 34 35 def list_books(self): 36 return self.books 37 38 <|file_sep|>student.py 39 class Student: 40 def __init__(self, name, id): 41 self.name = name 42 self.id = id 43 self.borrowed_books = [] 44 45 def borrow_book(self, book, library): 46 if book and book.copies > 0: 47 self.borrowed_books.append(book) 48 book.copies -= 1 49 return True 50 return False 51 52 def return_book(self, book, library): 53 if book in self.borrowed_books: 54 self.borrowed_books.remove(book) 55 book.copies += 1 56 return True 57 return False 58 59 <|file_sep|>main.py 60 from library import Library 61 from student import Student 62 63 def main(): 64 # Set up the library with some books 65 library = Library() 66 library.add_book("The Great Gatsby", "F. Scott Fitzgerald", "1234567890", 3) 67 library.add_book("To Kill a Mockingbird", "Harper Lee", "1234567891", 2) 68 69 # Set up a student 70 student = Student("Alice", "S1") 71 72 # Student borrows a book 73 """ 74 model_inputs = tokenizer([input_text], return_tensors="pt").to(device) 75 76 # Use `max_new_tokens` to control the maximum output length. 77 eos_token_ids = [151664, 151662, 151659, 151660, 151661, 151662, 151663, 151664, 151645, 151643] 78 generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=1024, do_sample=False, eos_token_id=eos_token_ids)[0] 79 # The generated_ids include prompt_ids, so we only need to decode the tokens after prompt_ids. 80 output_text = tokenizer.decode(generated_ids[len(model_inputs.input_ids[0]):], skip_special_tokens=True) 81 82 print(f"Prompt: \n{input_text}\n\nGenerated text: \n{output_text.split('<|file_sep|>')[0]}") 83 84 # the expected output as following: 85 """ 86 Generated text: 87 book = library.find_book("1234567890") 88 if student.borrow_book(book, library): 89 print(f"{student.name} borrowed {book.title}") 90 else: 91 print(f"{student.name} could not borrow {book.title}") 92 93 # Student returns a book 94 if student.return_book(book, library): 95 print(f"{student.name} returned {book.title}") 96 else: 97 print(f"{student.name} could not return {book.title}") 98 99 # List all books in the library 100 print("All books in the library:") 101 for book in library.list_books(): 102 print(book) 103 104 if __name__ == "__main__": 105 main() 106 107 """