Update app.py

app.py

@@ -1,38 +1,34 @@
 import os
 import gradio as gr
 from transformers import AutoTokenizer, pipeline
-import torch
 
-# Initialize the model and tokenizer
+# Initialize the model and tokenizer with environment variable for HF_TOKEN
 model_name = "AIFS/Prometh-MOEM-V.01"
+hf_token = os.getenv("HF_TOKEN")  # More Pythonic way to fetch environment variables
 
-HF_TOKEN = os.environ.get("HF_TOKEN")
-
-
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
 
 text_generation_pipeline = pipeline(
     "text-generation",
     model=model_name,
-    model_kwargs={"torch_dtype": torch.float16, "load_in_4bit": True},
-    use_auth_token=HF_TOKEN
+    model_kwargs={"torch_dtype": "auto", "load_in_4bit": True},  # 'auto' lets PyTorch decide the most optimal dtype
+    use_auth_token=hf_token
 )
 
-
 def generate_text(user_input):
     messages = [{"role": "user", "content": user_input}]
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     outputs = text_generation_pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
     return outputs[0]["generated_text"]
 
-# Create the Gradio interface
+# Updated Gradio interface creation to use the latest syntax
 iface = gr.Interface(
     fn=generate_text,
-    inputs=gr.inputs.Textbox(lines=2, placeholder="Type your question here..."),
-    outputs=gr.outputs.Textbox(),
+    inputs=gr.Textbox(lines=2, placeholder="Type your question here..."),
+    outputs=gr.Textbox(),
     title="Prometh-MOEM Text Generation",
     description="A text generation model that understands your queries and generates concise, informative responses."
 )
 
-# Run the interface
-iface.launch()
+# Run the interface with enhanced parameters for better performance and user experience
+iface.launch(enable_queue=True)  # enable_queue=True for handling high traffic