Update app.py

app.py

@@ -1,22 +1,30 @@
 import os
+import torch
+
 import streamlit as st
+
 from dotenv import load_dotenv
 from peft import PeftModel, PeftConfig
 from chromadb import HttpClient
 from utils.embedding_utils import CustomEmbeddingFunction
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-st.title("FormulAI Q&A")
+st.title("FormulAI")
+st.write("Benvenuto FormulaAI il Chatbot riguardante la Formula Uno! Chiedimi ciò che vuoi a riguardo!")
+st.write("I am a chatbot that has been fine-tuned on the FormuLLaMa-3.2-1B dataset.")
 
+# Device and model configuration
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model_name = "unsloth/Llama-3.2-1B"
 
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Load pretrained model and tokenizer
+model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
+# Load PEFT configuration and apply to model on device
 adapter_name = "FormulAI/FormuLLaMa-3.2-1B-LoRA"
 peft_config = PeftConfig.from_pretrained(adapter_name)
-
-model = PeftModel(model, peft_config)
+model = PeftModel(model, peft_config).to(device)
 
 template = """Answer the following QUESTION based on the CONTEXT given.
 If you do not know the answer and the CONTEXT doesn't contain the answer truthfully say "I don't know".
@@ -37,7 +45,7 @@ if 'past' not in st.session_state:
     st.session_state['past'] = []
 
 def get_text():
-    input_text = st.text_input("Ask your question: ", "", key="input")
+    input_text = st.text_input("Chiedi qualcosa: ", "", key="input")
     return input_text 
 
 load_dotenv("chroma.env")
@@ -57,12 +65,11 @@ if question:
     context = " ".join(response['documents'][0])
 
     input_text = template.replace("{context}", context).replace("{question}", question)
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
 
     output = model.generate(input_ids, max_new_tokens=200, early_stopping=True)
     answer = tokenizer.decode(output[0], skip_special_tokens=True).split("ANSWER:")[1]
 
-
     st.session_state.past.append(question)
     st.session_state.generated.append(answer)