ysn-rfd/text-dataset-tiny-code-script-py-format

1#!/usr/bin/env python
2import torch
3from transformers import AutoTokenizer, AutoModelForCausalLM
4
5# Define the directory where your fine-tuned model is saved.
6model_dir = "./gpt2-finetuned"
7
8# Load the tokenizer and model from the saved directory.
9tokenizer = AutoTokenizer.from_pretrained(model_dir)
10model = AutoModelForCausalLM.from_pretrained(model_dir)
11
12# If you are using GPU and it's available, move the model to GPU.
13device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14model.to(device)
15
16print("Chat with the model! Type 'exit' or 'quit' to end the conversation.")
17
18while True:
19    # Get user input.
20    user_input = input("You: ")
21    if user_input.lower() in ["exit", "quit"]:
22        print("Exiting chat.")
23        break
24
25    # Encode the input text and generate an attention mask.
26    inputs = tokenizer(user_input, return_tensors="pt", padding=True, truncation=True)
27    input_ids = inputs["input_ids"].to(device)
28    attention_mask = inputs["attention_mask"].to(device)  # Explicitly set the attention mask
29
30    # Generate a response. You can tweak the generation parameters as needed.
31    output_ids = model.generate(
32        input_ids,
33        attention_mask=attention_mask,  # Pass the attention mask here
34        max_length=100,             # Maximum length of the generated response.
35        do_sample=True,             # Use sampling; set to False for greedy decoding.
36        top_p=0.95,                 # Top-p (nucleus) sampling.
37        top_k=50,                   # Top-k sampling.
38        pad_token_id=tokenizer.eos_token_id  # Avoid warnings if no pad token is defined.
39    )
40
41    # Decode the generated tokens to a string.
42    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
43    print("Bot:", response)
44