dnnsdunca commited on
Commit
044208c
1 Parent(s): de5fa47

Create Agent-bert.py

Browse files
Files changed (1) hide show
  1. Agent-bert.py +40 -0
Agent-bert.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
3
+ from transformers import pipeline
4
+ model_name = "dbernsohn/roberta-java"
5
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
6
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
7
+ def preprocess_input(description):
8
+ input_text = "Generate an agent that " + description
9
+ inputs = tokenizer.encode(input_text, return_tensors='pt')
10
+ return inputs
11
+ def generate_agent_code(inputs):
12
+ generated_ids = model.generate(inputs)
13
+ agent_code = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
14
+ return agent_code
15
+ import os
16
+ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
17
+ from transformers import pipeline
18
+
19
+ # Load the pre-trained CodeBERTa model and tokenizer
20
+ model_name = "dbernsohn/roberta-java"
21
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
22
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
23
+
24
+ # Function to pre-process user input description
25
+ def preprocess_input(description):
26
+ input_text = "Generate an agent that " + description
27
+ inputs = tokenizer.encode(input_text, return_tensors='pt')
28
+ return inputs
29
+
30
+ # Function to generate agent code using the fine-tuned model
31
+ def generate_agent_code(inputs):
32
+ generated_ids = model.generate(inputs)
33
+ agent_code = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
34
+ return agent_code
35
+
36
+ # Example usage
37
+ user_description = "can perform sentiment analysis on text data."
38
+ inputs = preprocess_input(user_description)
39
+ generated_code = generate_agent_code(inputs)
40
+ print(generated_code)