Views
No views yet
1from transformers import T5Tokenizer, T5ForConditionalGeneration
2
3# Load pre-trained model and tokenizer
4model_name = 'VPrashant/cypher-gen'
5tokenizer = T5Tokenizer.from_pretrained(model_name)
6model = T5ForConditionalGeneration.from_pretrained(model_name)
7
8# Example input for testing
9test_input = "Which employees joined the company after 2015?"
10test_encoding = tokenizer(test_input, return_tensors="pt", max_length=128, truncation=True, padding="max_length")
11
12# Generate Cypher query
13output = model.generate(input_ids=test_encoding['input_ids'], max_length=128)
14generated_query = tokenizer.decode(output[0], skip_special_tokens=True)
15
16print("Generated Cypher Query:", generated_query)
17