File size: 1,140 Bytes
87c1d92 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
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() |