Rahul-8799 commited on
Commit
0b4887e
·
verified ·
1 Parent(s): 6e7849a

Create project_manager_agent.py

Browse files
Files changed (1) hide show
  1. agents/project_manager_agent.py +14 -0
agents/project_manager_agent.py ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoTokenizer, AutoModelForCausalLM
2
+ import torch
3
+
4
+ MODEL_REPO = "Rahul-8799/project_manager_gemma3"
5
+
6
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, trust_remote_code=True)
7
+ model = AutoModelForCausalLM.from_pretrained(MODEL_REPO, torch_dtype=torch.float16, device_map="auto")
8
+ model.eval()
9
+
10
+ def run(prompt):
11
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
12
+ with torch.no_grad():
13
+ outputs = model.generate(**inputs, max_new_tokens=512)
14
+ return tokenizer.decode(outputs[0], skip_special_tokens=True)