Update app.py

app.py

@@ -1,3 +1,36 @@
 import gradio as gr
 
-gr.load("models/mistralai/Mistral-7B-Instruct-v0.3").launch()
+gr.load("models/mistralai/Mistral-7B-Instruct-v0.3").launch()
+import gradio as gr
+import torch
+from mamba_model import MambaModel
+
+# Load the model
+model = MambaModel.from_pretrained(pretrained_model_name="Zyphra/BlackMamba-2.8B")
+model = model.cuda().half()
+
+# Define the function to generate output
+def generate_output(input_text):
+    # Convert the input text (comma-separated numbers) to a list of integers
+    try:
+        input_ids = [int(x.strip()) for x in input_text.split(",")]
+        inputs = torch.tensor(input_ids).cuda().long().unsqueeze(0)
+        
+        # Run the model and get output
+        with torch.no_grad():
+            out = model(inputs)
+        
+        # Convert output to a human-readable format
+        return out.cpu().numpy().tolist()
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Set up the Gradio interface
+input_component = gr.Textbox(label="Input IDs (comma-separated)", placeholder="Enter input IDs like: 1, 2")
+output_component = gr.Textbox(label="Output")
+
+iface = gr.Interface(fn=generate_output, inputs=input_component, outputs=output_component, title="BlackMamba Model")
+
+# Launch the interface
+if __name__ == "__main__":
+    iface.launch()