Views
No views yet
1from transformers import T5ForConditionalGeneration, T5Tokenizer
2device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
3
4# Load FP16 model
5model = T5ForConditionalGeneration.from_pretrained("./t5_email_finetuned_fp16").to(device)
6tokenizer = T5Tokenizer.from_pretrained("./t5_email_finetuned_fp16")
7
8# Generate a response
9def generate_response(prompt, max_length=128):
10 input_text = f"generate response: {prompt}"
11 inputs = tokenizer(input_text, max_length=128, truncation=True, padding="max_length", return_tensors="pt").to(device)
12 outputs = model.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], max_length=max_length, num_beams=4, early_stopping=True)
13 return tokenizer.decode(outputs[0], skip_special_tokens=True)
14
15# Example
16print(generate_response("Can you send me the report?"))