Spaces:
Sleeping
Sleeping
File size: 1,701 Bytes
b092158 884a117 b092158 1be0960 b092158 884a117 d248649 884a117 d248649 884a117 0f8a863 884a117 befa9fa 091b809 b092158 c862f71 091b809 884a117 befa9fa 47da2b4 befa9fa 47da2b4 091b809 b092158 4347369 b092158 | 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 | import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import multiprocessing
model_id = "distilgpt2"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
def run_generation(prompt, return_dict):
try:
output = generator(
prompt,
max_new_tokens=48,
do_sample=True,
temperature=0.7,
eos_token_id=tokenizer.eos_token_id
)[0]["generated_text"]
return_dict["result"] = output
except Exception as e:
return_dict["result"] = f"GENERATION ERROR: {e}"
def generate_with_hard_timeout(prompt, timeout=15):
manager = multiprocessing.Manager()
return_dict = manager.dict()
p = multiprocessing.Process(target=run_generation, args=(prompt, return_dict))
p.start()
p.join(timeout)
if p.is_alive():
p.terminate()
return "ERROR: Generation timed out."
return return_dict["result"]
def chat(input_text):
prompt = input_text + "\n"
try:
output = generate_with_hard_timeout(prompt)
if isinstance(output, str):
reply = output[len(prompt):].strip()
if not reply or reply.isspace():
reply = output.strip()
else:
reply = "ERROR: Unexpected output format."
return reply
except Exception as e:
return f"GENERATION ERROR: {e}"
demo = gr.Interface(
fn=chat,
inputs=gr.Textbox(label="input_text"),
outputs="text",
title="Kairon (Unprimed)",
allow_flagging="never"
)
demo.queue()
demo.launch()
|