Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline | |
tokenizer = AutoTokenizer.from_pretrained(".", local_files_only=True) | |
model = AutoModelForSeq2SeqLM.from_pretrained(".", local_files_only=True) | |
pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer) | |
def predict(task, prompt, context="", auto_cot=False): | |
full_prompt = f"[TASK: {task.upper()}] {prompt}" | |
if context: | |
full_prompt += f" Context: {context}" | |
if auto_cot: | |
full_prompt += "\nLet's think step by step." | |
output = pipe(full_prompt, max_new_tokens=128)[0]["generated_text"] | |
return output | |
demo = gr.Interface( | |
fn=predict, | |
inputs=["text", "text", "text", "checkbox"], | |
outputs="text" | |
) | |
demo.launch() | |