MAmmoTH2 / app.py
aaabiao's picture
Update app.py
d0f1436 verified
raw
history blame
545 Bytes
import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("TIGER-Lab/MAmmoTH2-8B-Plus", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("TIGER-Lab/MAmmoTH2-8B-Plus", trust_remote_code=True)
def predict(input_text):
outputs = model.generate(input_text, max_length=50, do_sample=True)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
demo = gr.Interface(fn=predict, inputs="text", outputs="text")
demo.launch()