File size: 1,859 Bytes
55f9709 a5697da 55f9709 0051fda 55f9709 0051fda 55f9709 |
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 pipeline
import torch
MAX_NEW_TOKENS = 250
MODEL="HuggingFaceTB/SmolLM2-135M-Instruct"
# MODEL="HuggingFaceTB/SmolLM2-360M-Instruct"
# MODEL="HuggingFaceTB/SmolLM2-1.7B-Instruct"
TEMPERATURE = 0.6
TOP_P = 0.95
REPETITION_PENALTY = 1.2
pipe = pipeline("text-generation", model="HuggingFaceTB/SmolLM2-1.7B-Instruct")
def message_fx(message, history):
if len(history) == 0:
send_to_api = [{'role':'user', 'content':message}]
print(send_to_api)
with torch.no_grad():
response = pipe(send_to_api,
do_sample=True,
max_new_tokens=MAX_NEW_TOKENS,
temperature=TEMPERATURE, # 1.0 = lots of creativity, high odd of hallucination 0.1 very specific writing and low odds
# top_k=50,
top_p=TOP_P,
repetition_penalty=REPETITION_PENALTY, # Added to discourage repetition
# no_repeat_ngram_size=3
)[0]['generated_text'][1]['content']
return response
else:
send_to_api = history + [{'role':'user', 'content':message}]
print(send_to_api)
with torch.no_grad():
response = pipe(send_to_api,
do_sample=True,
max_new_tokens=MAX_NEW_TOKENS,
temperature=TEMPERATURE, # 1.0 = lots of creativity, high odd of hallucination 0.1 very specific writing and low odds
# top_k=50,
top_p=TOP_P,
repetition_penalty=REPETITION_PENALTY, # Added to discourage repetition
# no_repeat_ngram_size=3
)[0]['generated_text'][-1]['content']
return response
gr.ChatInterface(
fn=message_fx,
type="messages"
).launch()
|