File size: 1,014 Bytes
fc9d193
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load your fine-tuned model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("./lockin_model")
model = AutoModelForCausalLM.from_pretrained("./lockin_model")

# Function to generate yes/no questions
def generate_question(input_text):
    # Add padding and attention mask
    inputs = tokenizer(
        input_text,
        return_tensors="pt",
        padding=True,
        truncation=True,
        return_attention_mask=True
    )
    
    output = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],  # Add attention mask
        max_new_tokens=100,
        do_sample=True,
        temperature=1.5,
        top_p=0.8,
        top_k=50,
        pad_token_id=tokenizer.eos_token_id  # Explicitly set pad token ID
    )
    return tokenizer.decode(output[0], skip_special_tokens=True)

# Example usage
prompt = "What the fuck"
question = generate_question(prompt)
print("Generated Question:", question)