Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| import torch | |
| import os | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| api_key = os.getenv("api_key") | |
| # App title and description | |
| st.title("I am Your GrowBuddy 🌱") | |
| st.write("Let me help you start gardening. Let's grow together!") | |
| def load_model(): | |
| try: | |
| tokenizer = AutoTokenizer.from_pretrained("KhunPop/Gardening", use_auth_token=api_key) | |
| model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", use_auth_token=api_key) | |
| return tokenizer, model | |
| except Exception as e: | |
| st.error(f"Failed to load model: {e}") | |
| return None, None | |
| # Load model and tokenizer | |
| tokenizer, model = load_model() | |
| if not tokenizer or not model: | |
| st.stop() | |
| # Default to CPU, or use GPU if available | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = model.to(device) | |
| # Initialize session state messages if not already initialized | |
| if "messages" not in st.session_state: | |
| st.session_state.messages = [ | |
| {"role": "assistant", "content": "Hello there! How can I help you with gardening today?"} | |
| ] | |
| # Display the conversation history | |
| for message in st.session_state.messages: | |
| with st.chat_message(message["role"]): | |
| st.write(message["content"]) | |
| def generate_response(prompt): | |
| try: | |
| # Tokenize the input prompt | |
| inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device) | |
| # Ensure the model is generating properly (without a target) | |
| outputs = model.generate(inputs["input_ids"], max_new_tokens=150, temperature=0.7, do_sample=True) | |
| # Decode the output to text | |
| response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| return response | |
| except Exception as e: | |
| st.error(f"Error during text generation: {e}") | |
| return "Sorry, I couldn't process your request." | |
| # User input field for asking questions | |
| user_input = st.chat_input("Type your gardening question here:") | |
| if user_input: | |
| # Display user message | |
| with st.chat_message("user"): | |
| st.write(user_input) | |
| # Generate and display assistant's response | |
| with st.chat_message("assistant"): | |
| with st.spinner("I'm gonna tell you..."): | |
| response = generate_response(user_input) | |
| st.write(response) | |
| # Update session state with the new conversation | |
| st.session_state.messages.append({"role": "user", "content": user_input}) | |
| st.session_state.messages.append({"role": "assistant", "content": response}) | |