Update app.py

app.py

@@ -1,25 +1,30 @@
-from transformers import GemmaForCausalLM, AutoTokenizer
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load tokenizer
+# Load the model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("umar141/Gemma_1B_Baro_v2_vllm")
+model = AutoModelForCausalLM.from_pretrained("umar141/Gemma_1B_Baro_v2_vllm")
 
-# Load model
-model = GemmaForCausalLM.from_pretrained(
-    "umar141/Gemma_1B_Baro_v2_vllm",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto"
-)
+# Streamlit page configuration
+st.set_page_config(page_title="Gemma-based Chatbot", page_icon=":robot:")
 
-# Tokenize prompt
-input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+# Introduction text
+st.title("Gemma-based Chatbot")
+st.write("This is a chatbot powered by a fine-tuned Gemma model.")
 
-# Generate
-outputs = model.generate(
-    input_ids=input_ids,
-    max_new_tokens=200,
-    do_sample=True,
-    top_p=0.9,
-    temperature=0.7,
-)
+# User input
+user_input = st.text_input("Ask me anything:")
+
+# Generate response when the user inputs a query
+if user_input:
+    # Tokenize input and generate model response
+    inputs = tokenizer.encode(user_input, return_tensors="pt")
+    outputs = model.generate(inputs, max_length=150, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
+    
+    # Decode the response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Display the response
+    st.write("AI Response:")
+    st.write(response)
 
-response = tokenizer.decode(outputs[0], skip_special_tokens=True)