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