App.py

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()