Spaces:
Runtime error
Runtime error
from transformers import pipeline | |
from transformers import StoppingCriteria, StoppingCriteriaList | |
from transformers import AutoTokenizer | |
import torch | |
tokenizer = AutoTokenizer.from_pretrained("Hieu-Pham/Llama2-7B-QLoRA-cooking-text-gen-merged") | |
stop_token_ids = tokenizer.convert_tokens_to_ids(["\n", "#", "\\", "`", "###", "##", "Question", "Comment", "Answer", "Context"]) | |
# define custom stopping criteria object | |
class StopOnTokens(StoppingCriteria): | |
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: | |
for stop_id in stop_token_ids: | |
if input_ids[0][-1] == stop_id: | |
return True | |
return False | |
stopping_criteria = StoppingCriteriaList([StopOnTokens()]) | |
pipe = pipeline( | |
task="text-generation", | |
model="Hieu-Pham/Llama2-7B-QLoRA-cooking-text-gen-merged", | |
tokenizer=tokenizer, | |
return_full_text=False, | |
stopping_criteria=stopping_criteria, | |
temperature=0.1, | |
top_p=0.15, | |
top_k=0, | |
max_new_tokens=100, | |
repetition_penalty=1.1 | |
) | |
import gradio as gr | |
def predict(question, context): | |
input = f"Question: {question} Context: {context} Answer:" | |
predictions = pipe(input) | |
output = predictions[0]["generated_text"].replace("Question", "") | |
output = output.replace("Answer", "") | |
return output | |
demo = gr.Interface( | |
predict, | |
inputs=[gr.Textbox(lines=2, placeholder="Please provide your question", label="Question"), gr.Textbox(lines=2, placeholder="Please provide your context", label="Context")], | |
outputs=gr.Textbox(lines=2, placeholder="Answers...", label="Output"), | |
title="Cooking Recipe MRC", | |
) | |
demo.launch() |