Rahul-8799 commited on
Commit
8afb084
·
verified ·
1 Parent(s): 0ee2ec4

Delete product_manager_agent.py

Browse files
Files changed (1) hide show
  1. product_manager_agent.py +0 -27
product_manager_agent.py DELETED
@@ -1,27 +0,0 @@
1
- from transformers import AutoTokenizer, AutoModelForCausalLM
2
- import torch
3
- from langchain_core.messages import AIMessage
4
-
5
- MODEL_REPO = "Rahul-8799/product_manager_mistral"
6
-
7
- tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, trust_remote_code=True)
8
- model = AutoModelForCausalLM.from_pretrained(
9
- MODEL_REPO,
10
- torch_dtype=torch.float16,
11
- device_map="auto"
12
- )
13
-
14
- def run(state: dict) -> dict:
15
- """Generates structured product requirements from user input prompt."""
16
- messages = state["messages"]
17
- prompt = messages[-1].content
18
-
19
- input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
20
- output_ids = model.generate(input_ids, max_new_tokens=3000)
21
- output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
22
-
23
- return {
24
- "messages": [AIMessage(content=output)],
25
- "chat_log": state["chat_log"] + [{"role": "Product Manager", "content": output}],
26
- "pm_output": output,
27
- }