jayanth7111's picture
Update app.py
9d2acd8 verified
raw
history blame
765 Bytes
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()