Zoe911's picture
Update app.py
c4b0502 verified
raw
history blame contribute delete
698 Bytes
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
model_id = "NousResearch/Llama-2-7b-chat-hf" #
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
def extract_event(text):
prompt = f"""Extract the following fields from the input text:
- Event title
- Date (YYYY-MM-DD)
- Start Time (HH:MM)
- End Time (HH:MM)
- Duration
- Location
- People
Text: {text}
Return in JSON format.
"""
response = pipe(prompt, max_new_tokens=300, do_sample=False)[0]['generated_text']
return response
gr.Interface