Update app.py

app.py

@@ -1,10 +1,12 @@
+import os
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import os
+
+access_token = os.getenv('HF_TOKEN')
 
 # Define the repository ID and access token
 repo_id = "Mikhil-jivus/Llama-32-3B-FineTuned"
-access_token = os.getenv('HF_TOKEN')
+access_token = "your_access_token_here"
 
 # Load the tokenizer and model from the Hugging Face repository
 tokenizer = AutoTokenizer.from_pretrained(repo_id, use_auth_token=access_token)
@@ -32,6 +34,9 @@ def respond(
     input_text = system_message + " ".join([f"{msg['role']}: {msg['content']}" for msg in messages])
     input_ids = tokenizer.encode(input_text, return_tensors="pt")
 
+    # Create attention mask
+    attention_mask = input_ids.ne(tokenizer.pad_token_id).long()
+
     # Generate a response
     chat_history_ids = model.generate(
         input_ids,
@@ -40,6 +45,7 @@ def respond(
         top_p=top_p,
         pad_token_id=tokenizer.eos_token_id,
         do_sample=True,
+        attention_mask=attention_mask,
     )
 
     # Decode the response
@@ -67,4 +73,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)