Update app.py

app.py

@@ -4,21 +4,25 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import gradio as gr
 
-# Set token securely
+# Get token from environment (secure via Repository Secrets in HF Space)
 hf_token = os.getenv("HF_TOKEN")
 if not hf_token:
     raise ValueError("Please set your HF_TOKEN environment variable securely.")
 
+# Login to Hugging Face Hub
 login(token=hf_token)
 
+# Model and device setup
 model_name = "mistralai/Mistral-7B-Instruct-v0.2"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
+# Load model
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     device_map="auto" if torch.cuda.is_available() else None,
@@ -30,6 +34,7 @@ model = AutoModelForCausalLM.from_pretrained(
 if not torch.cuda.is_available():
     model = model.to(device)
 
+# Generation function
 def generate_text(prompt):
     if not prompt.strip():
         return "Please enter a valid prompt."
@@ -41,7 +46,14 @@ def generate_text(prompt):
             tokenize=False,
             add_generation_prompt=True
         )
-        inputs = tokenizer(formatted_prompt, return_tensors="pt", padding=True, truncation=True, max_length=2048)
+        
+        inputs = tokenizer(
+            formatted_prompt,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=2048
+        )
         inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
         with torch.no_grad():
@@ -54,7 +66,7 @@ def generate_text(prompt):
                 pad_token_id=tokenizer.eos_token_id,
                 eos_token_id=tokenizer.eos_token_id
             )
-
+        
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         generated_text = response[len(formatted_prompt):].strip()
         return generated_text
@@ -62,10 +74,18 @@ def generate_text(prompt):
     except Exception as e:
         return f"Error generating text: {str(e)}"
 
+# Gradio interface
 iface = gr.Interface(
     fn=generate_text,
-    inputs=gr.Textbox(lines=3, placeholder="Enter your prompt here...", label="Your Message"),
-    outputs=gr.Textbox(label="Mistral Response", lines=5),
+    inputs=gr.Textbox(
+        lines=3, 
+        placeholder="Enter your prompt here...",
+        label="Your Message"
+    ),
+    outputs=gr.Textbox(
+        label="Mistral Response",
+        lines=5
+    ),
     title="🤖 Mistral-7B Chat",
     description="Chat with Mistral-7B-Instruct model. Enter your message and get AI-generated responses.",
     examples=[
@@ -77,5 +97,10 @@ iface = gr.Interface(
     cache_examples=False
 )
 
+# Run app
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    iface.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )