Harish2002 commited on
Commit
ba67f04
·
verified ·
1 Parent(s): 9a9a960

Upload agent.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. agent.py +58 -0
agent.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import sys
2
+ import json
3
+ import os
4
+ from datetime import datetime
5
+ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
6
+ from peft import PeftModel, PeftConfig
7
+
8
+ # Load model and tokenizer
9
+ def load_model():
10
+ base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
11
+ adapter_path = "Harish2002/cli-lora-tinyllama" # ✅ fixed path
12
+
13
+ tokenizer = AutoTokenizer.from_pretrained(base_model)
14
+ model = AutoModelForCausalLM.from_pretrained(base_model)
15
+ model = PeftModel.from_pretrained(model, adapter_path)
16
+ return tokenizer, model
17
+
18
+ # Generate plan from input instruction
19
+ def generate_plan(prompt, tokenizer, model):
20
+ pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)
21
+ output = pipe(prompt)[0]['generated_text']
22
+ return output.strip()
23
+
24
+ # Check if first line is a shell command
25
+ def is_shell_command(line):
26
+ return line.startswith(("git", "bash", "tar", "gzip", "grep", "python", "./", "cd", "ls"))
27
+
28
+ # Log to logs/trace.jsonl
29
+ def log_trace(prompt, response):
30
+ os.makedirs("logs", exist_ok=True)
31
+ trace = {
32
+ "timestamp": datetime.utcnow().isoformat(),
33
+ "input": prompt,
34
+ "response": response
35
+ }
36
+ with open("logs/trace.jsonl", "a") as f:
37
+ f.write(json.dumps(trace) + "\n")
38
+
39
+ # Main
40
+ if __name__ == "__main__":
41
+ if len(sys.argv) < 2:
42
+ print("Usage: python agent.py \"Your instruction here\"")
43
+ sys.exit(1)
44
+
45
+ user_input = sys.argv[1]
46
+ tokenizer, model = load_model()
47
+ result = generate_plan(user_input, tokenizer, model)
48
+
49
+ # Print result and echo dry-run if it's a shell command
50
+ print("\nGenerated Plan:\n")
51
+ print(result)
52
+
53
+ first_line = result.splitlines()[0]
54
+ if is_shell_command(first_line):
55
+ print("\nDry-run:")
56
+ print(f"echo {first_line}")
57
+
58
+ log_trace(user_input, result)