Spaces:
Sleeping
Sleeping
File size: 2,177 Bytes
60a5d4d 995a873 60a5d4d 995a873 60a5d4d 995a873 60a5d4d 995a873 60a5d4d 995a873 60a5d4d 995a873 60a5d4d 995a873 60a5d4d 995a873 60a5d4d 995a873 60a5d4d 995a873 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import time
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
# Load model
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2")
# Inference function
def chat_completion(messages, model_name="mock-gpt-model", max_tokens=512, temperature=0.1, stream=False):
if not messages:
return {
"error": "No messages provided."
}
# Rebuild prompt
prompt = ""
for msg in messages:
role = msg.get("role", "")
content = msg.get("content", "")
if role == "user":
prompt += f"User: {content}\n"
elif role == "assistant":
prompt += f"Assistant: {content}\n"
prompt += "Assistant:"
# Generate output
inputs = tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=max_tokens,
temperature=temperature,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Extract assistant reply
assistant_reply = generated_text[len(prompt):].strip()
return {
"id": "1337",
"object": "chat.completion",
"created": time.time(),
"model": model_name,
"choices": [{
"message": {
"role": "assistant",
"content": assistant_reply
}
}]
}
# Gradio API endpoint setup
demo = gr.Interface(
fn=chat_completion,
inputs=[
gr.JSON(label="messages"), # List[{"role":..., "content":...}]
gr.Textbox(label="model", value="mock-gpt-model"),
gr.Slider(minimum=1, maximum=1024, value=512, label="max_tokens"),
gr.Slider(minimum=0.0, maximum=1.0, value=0.1, label="temperature"),
gr.Checkbox(label="stream", value=False)
],
outputs=gr.JSON(label="response"),
title="OpenAI-compatible Chat API (Gradio + Transformers)",
allow_flagging="never"
)
if __name__ == "__main__":
demo.launch()
|