Update app.py

app.py

@@ -1,18 +1,41 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
+import gradio as gr
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Read Hugging Face token from environment variable
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-model_name = "hari7261/TechChat"
+MODEL_NAME = "hari7261/TechChat"
 
-# If the repo is private, add your token:
-token = "your_huggingface_token_here"  # or set HF_TOKEN env var
+print("Loading tokenizer and model...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, use_auth_token=HF_TOKEN)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)
-model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=token)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
-prompt = "Hello, how can I help you today?"
-inputs = tokenizer(prompt, return_tensors="pt")
+def generate_text(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(
+        **inputs,
+        max_length=150,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+        eos_token_id=tokenizer.eos_token_id
+    )
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return text
 
-with torch.no_grad():
-    outputs = model.generate(**inputs, max_length=50)
+# Gradio UI
+iface = gr.Interface(
+    fn=generate_text,
+    inputs=gr.Textbox(lines=5, placeholder="Enter your prompt here..."),
+    outputs="text",
+    title="TechChat - Mistral 7B",
+    description="Generate text with hari7261/TechChat model hosted on Hugging Face."
+)
 
-print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+if __name__ == "__main__":
+    iface.launch()