Update app.py

app.py

@@ -1,3 +1,37 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Load NVLM-D-72B model and tokenizer
+# model_name = "nvidia/NVLM-D-72B"
+model_name = "nvidia/NVLM-D-7B"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    trust_remote_code=True, 
+    device_map="auto"
+)
+
+# Inference function
+def generate_response(prompt, max_tokens=50):
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")  # Adjust to "cpu" if GPU unavailable
+    outputs = model.generate(**inputs, max_new_tokens=max_tokens)
+    return tokenizer.decode(outputs[0])
+
+# Gradio interface
+interface = gr.Interface(
+    fn=generate_response,
+    inputs=[
+        gr.Textbox(lines=2, label="Enter your prompt"),
+        gr.Slider(10, 100, step=10, value=50, label="Max Tokens")
+    ],
+    outputs="text",
+    title="NVIDIA NVLM-D-72B Demo",
+    description="Generate text using NVIDIA's NVLM-D-72B model."
+)
+if __name__ == "__main__":
+    interface.launch()
+
+
 # import gradio as gr
 # from huggingface_hub import InferenceClient
 
@@ -63,34 +97,3 @@
 # if __name__ == "__main__":
 #     demo.launch()
 
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-# Load NVLM-D-72B model and tokenizer
-model_name = "nvidia/NVLM-D-72B"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name, 
-    trust_remote_code=True, 
-    device_map="auto"
-)
-
-# Inference function
-def generate_response(prompt, max_tokens=50):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")  # Adjust to "cpu" if GPU unavailable
-    outputs = model.generate(**inputs, max_new_tokens=max_tokens)
-    return tokenizer.decode(outputs[0])
-
-# Gradio interface
-interface = gr.Interface(
-    fn=generate_response,
-    inputs=[
-        gr.Textbox(lines=2, label="Enter your prompt"),
-        gr.Slider(10, 100, step=10, value=50, label="Max Tokens")
-    ],
-    outputs="text",
-    title="NVIDIA NVLM-D-72B Demo",
-    description="Generate text using NVIDIA's NVLM-D-72B model."
-)
-if __name__ == "__main__":
-    interface.launch()