JanviMl commited on
Commit
b9f8cf8
·
verified ·
1 Parent(s): 963312c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +94 -0
app.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ import torch
4
+ import firebase_admin
5
+ from firebase_admin import credentials, db
6
+ import os
7
+ import json
8
+
9
+ # Load Firebase credentials from firebase-key.json
10
+ firebase_key_path = os.environ.get("FIREBASE_KEY_PATH", "firebase-key.json")
11
+ with open(firebase_key_path, "r") as f:
12
+ firebase_config = json.load(f)
13
+
14
+ # Initialize Firebase
15
+ cred = credentials.Certificate(firebase_config)
16
+ firebase_admin.initialize_app(cred, {
17
+ "databaseURL": "https://taskmate-d6e71-default-rtdb.firebaseio.com/" # Confirm this URL!
18
+ })
19
+ ref = db.reference("tasks")
20
+
21
+ # Load IBM Granite model from Hugging Face
22
+ model_name = "ibm-granite/granite-7b-base—" # Switch to "granite-3b" if 7b is too heavy
23
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
24
+ model = AutoModelForCausalLM.from_pretrained(model_name)
25
+
26
+ # Function to generate text with Granite
27
+ def generate_response(prompt, max_length=100):
28
+ inputs = tokenizer(prompt, return_tensors="pt")
29
+ outputs = model.generate(**inputs, max_length=max_length, num_return_sequences=1)
30
+ return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
31
+
32
+ # Parse user input into structured task
33
+ def parse_task(input_text, persona="default"):
34
+ prompt = f"For a {persona} employee, extract task, time, priority from: '{input_text}'"
35
+ response = generate_response(prompt)
36
+ return response # e.g., "Task: Email boss, Time: Today, Priority: High"
37
+
38
+ # Generate persona-specific subtasks
39
+ def generate_subtasks(task, persona="default"):
40
+ prompt = f"List 3 subtasks for '{task}' suited for a {persona} employee."
41
+ response = generate_response(prompt, max_length=150)
42
+ return response # e.g., "1. Draft email\n2. Send it\n3. Chill"
43
+
44
+ # Main chat function
45
+ def task_mate_chat(user_input, persona, chat_history):
46
+ # Parse the input
47
+ parsed = parse_task(user_input, persona)
48
+ task_name = parsed.split(",")[0].replace("Task: ", "").strip()
49
+
50
+ # Generate subtasks
51
+ subtasks = generate_subtasks(task_name, persona)
52
+
53
+ # Store in Firebase
54
+ task_data = {
55
+ "input": user_input,
56
+ "parsed": parsed,
57
+ "subtasks": subtasks,
58
+ "persona": persona,
59
+ "timestamp": str(db.ServerValue.TIMESTAMP)
60
+ }
61
+ ref.push().set(task_data)
62
+
63
+ # Format response
64
+ response = f"Parsed: {parsed}\nSubtasks:\n{subtasks}"
65
+ chat_history.append((user_input, response))
66
+ return "", chat_history
67
+
68
+ # Gradio Interface
69
+ with gr.Blocks(title="Task_Mate") as interface:
70
+ gr.Markdown("# Task_Mate: Your AI Task Buddy")
71
+ persona = gr.Dropdown(["lazy", "multitasker", "perfect"], label="Who are you?", value="lazy")
72
+ chatbot = gr.Chatbot(label="Chat with Task_Mate")
73
+ msg = gr.Textbox(label="Talk to me", placeholder="e.g., 'What’s today?' or 'Meeting at 2 PM'")
74
+ submit = gr.Button("Submit")
75
+
76
+ # Handle chat submission
77
+ submit.click(
78
+ fn=task_mate_chat,
79
+ inputs=[msg, persona, chatbot],
80
+ outputs=[msg, chatbot]
81
+ )
82
+
83
+ # Examples for each persona
84
+ gr.Examples(
85
+ examples=[
86
+ ["What’s today?", "lazy"],
87
+ ["Meeting Sarah, slides, IT call", "multitasker"],
88
+ ["Email boss by 3 PM", "perfect"]
89
+ ],
90
+ inputs=[msg, persona],
91
+ outputs=chatbot
92
+ )
93
+
94
+ interface.launch()