app.py

import gradio as gr
from huggingface_hub import InferenceClient
from optimum.intel import OVModelForCausalLM
from transformers import AutoTokenizer, pipeline
import re

# 載入模型和標記器
model_id = "hsuwill000/SmolLM2-135M-openvino"
model = OVModelForCausalLM.from_pretrained(model_id, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_id)

# 建立生成管道
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

def remove_incomplete_sentence(text):
    # 將文章根據句點、問號或感嘆號來分割成句子
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    
    # 檢查最後一句是否以省略號結尾或沒有標點符號結尾
    last_sentence = sentences[-1]
    if last_sentence.endswith('...') or not re.match(r'.*[.!?]$', last_sentence):
        # 移除最後一句不完整句子
        sentences = sentences[:-1]
    
    # 返回重新組合的文章
    return ' '.join(sentences)

def respond(message, history):
    # 將當前訊息與歷史訊息合併
    input_text = message if not history else history[-1]["content"] + " " + message
    input_text = message
    # 獲取模型的回應
    response = pipe(input_text, max_new_tokens=150, truncation=True, num_return_sequences=1)
    reply = remove_incomplete_sentence ( response[0]['generated_text'] )
    
    # 返回新的消息格式
    print(f"Message: {message}")
    print(f"Reply: {reply}")
    return reply
    
# 設定 Gradio 的聊天界面
demo = gr.ChatInterface(fn=respond, title="Chat with SmolLM2-135M-openvino(only English):", description="SmolLM2-135M-openvino", type='messages')

if __name__ == "__main__":
    demo.launch()