Mattral commited on
Commit
3d0f7c4
·
verified ·
1 Parent(s): df39778

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +69 -0
app.py ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ import torch
4
+ from PIL import Image
5
+ import os
6
+
7
+
8
+ from dotenv import load_dotenv
9
+ load_dotenv()
10
+
11
+ if not api_key:
12
+ st.error("Hugging Face API key is not set. Please check your .env file or environment variables.")
13
+
14
+ # App title and description
15
+ st.title("I am Your GrowBuddy 🌱")
16
+ st.write("Let me help you start gardening. Let's grow together!")
17
+
18
+ def load_model():
19
+ try:
20
+ tokenizer = AutoTokenizer.from_pretrained("KhunPop/Gardening")
21
+ model = AutoModelForCausalLM.from_pretrained("unsloth/gemma-2-2b")
22
+ return tokenizer, model
23
+ except Exception as e:
24
+ st.error(f"Failed to load model: {e}")
25
+ return None, None
26
+
27
+ tokenizer, model = load_model()
28
+
29
+ if not tokenizer or not model:
30
+ st.stop()
31
+
32
+ # Default to CPU
33
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
34
+ if model is not None:
35
+ model = model.to(device)
36
+
37
+ if "messages" not in st.session_state.keys():
38
+ st.session_state.messages = [
39
+ {"role": "assistant", "content": "Hello there! How can I help you with gardening today?"}
40
+ ]
41
+
42
+ for message in st.session_state.messages:
43
+ with st.chat_message(message["role"]):
44
+ st.write(message["content"])
45
+
46
+ def generate_response(prompt):
47
+ try:
48
+ inputs = tokenizer(prompt, return_tensors="pt").to(device) # Ensure inputs are moved to the device (CPU)
49
+ outputs = model.generate(**inputs, max_new_tokens=150, temperature=0.7)
50
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
51
+ return response
52
+ except Exception as e:
53
+ st.error(f"Error during text generation: {e}")
54
+ return "Sorry, I couldn't process your request."
55
+
56
+ user_input = st.chat_input("Type your gardening question here:")
57
+
58
+ if user_input:
59
+ # Display user message
60
+ with st.chat_message("user"):
61
+ st.write(user_input)
62
+
63
+ with st.chat_message("assistant"):
64
+ with st.spinner("I'm gonna tell you..."):
65
+ response = generate_response(user_input)
66
+ st.write(response)
67
+
68
+ st.session_state.messages.append({"role": "user", "content": user_input})
69
+ st.session_state.messages.append({"role": "assistant", "content": response})