/ examples / Qwen2.5-Coder-repolevel.py
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  """