Spaces:
Sleeping
Sleeping
File size: 2,604 Bytes
3d0f7c4 abfd83b 5d4ec37 3d0f7c4 5d4ec37 3d0f7c4 5d4ec37 3d0f7c4 ddcad02 3d0f7c4 ddcad02 3d0f7c4 ddcad02 3d0f7c4 ddcad02 3d0f7c4 791b656 ddcad02 791b656 3d0f7c4 ddcad02 3d0f7c4 ddcad02 3d0f7c4 ddcad02 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
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})
|