phi-2-demo / app.py
sombochea's picture
Add phi-2
4d214ec unverified
raw
history blame contribute delete
905 Bytes
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
# Load model and tokenizer (using CPU for broader accessibility)
model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", torch_dtype=torch.float32, device_map="cpu", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
def generate_text(prompt):
inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False)
outputs = model.generate(**inputs, max_length=200)
text = tokenizer.batch_decode(outputs)[0]
return text
# Create Gradio interface
iface = gr.Interface(
fn=generate_text,
inputs=[gr.Textbox(lines=5, label="Enter your prompt")],
outputs="text",
title="PHI-2 Text Generator",
description="Generate text using the PHI-2 generative language model",
)
# Launch the interface
iface.launch()