|
import torch |
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
import gradio as gr |
|
import traceback |
|
|
|
|
|
model_name = "lei-HuggingFace/SFT_Level_Measurement_Guide_07182024" |
|
|
|
def load_model(): |
|
try: |
|
model = AutoModelForCausalLM.from_pretrained( |
|
model_name, |
|
device_map="auto", |
|
torch_dtype=torch.float16, |
|
trust_remote_code=True |
|
) |
|
print(f"Model loaded successfully. Device: {model.device}") |
|
return model |
|
except Exception as e: |
|
print(f"Error loading model: {str(e)}") |
|
traceback.print_exc() |
|
return None |
|
|
|
try: |
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
print("Tokenizer loaded successfully") |
|
except Exception as e: |
|
print(f"Error loading tokenizer: {str(e)}") |
|
traceback.print_exc() |
|
tokenizer = None |
|
|
|
model = load_model() |
|
|
|
def generate_response(message, history): |
|
try: |
|
|
|
messages = [] |
|
for h in history: |
|
messages.append({"role": "user", "content": h[0]}) |
|
messages.append({"role": "assistant", "content": h[1]}) |
|
messages.append({"role": "user", "content": message}) |
|
|
|
print(f"Prepared messages: {messages}") |
|
|
|
|
|
input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device) |
|
print(f"Input shape: {input_ids.shape}") |
|
|
|
|
|
with torch.no_grad(): |
|
output = model.generate( |
|
input_ids, |
|
max_new_tokens=500, |
|
do_sample=True, |
|
temperature=0.7, |
|
pad_token_id=tokenizer.eos_token_id |
|
) |
|
|
|
response = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True) |
|
print(f"Generated response: {response}") |
|
return response.strip() |
|
except Exception as e: |
|
error_message = f"Error generating response: {str(e)}" |
|
print(error_message) |
|
traceback.print_exc() |
|
return error_message |
|
|
|
|
|
iface = gr.ChatInterface( |
|
generate_response, |
|
chatbot=gr.Chatbot(height=300), |
|
textbox=gr.Textbox(placeholder="Type your message here...", container=False, scale=7), |
|
title="Level Measurement Guide Chatbot", |
|
description="Chat with the fine-tuned Level Measurement Guide model.", |
|
theme="soft", |
|
examples=[ |
|
"What are the key considerations for level measurement in industrial settings?", |
|
"Can you explain the principle behind ultrasonic level sensors?", |
|
"What are the advantages of using radar level sensors?", |
|
], |
|
cache_examples=False, |
|
retry_btn=None, |
|
undo_btn="Delete Previous", |
|
clear_btn="Clear", |
|
) |
|
|
|
|
|
iface.launch(share=True, debug=True) |