75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
import os
|
|
import torch
|
|
import gradio as gr
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
import time
|
|
|
|
# Identity and Developer info
|
|
DEVELOPER = "Mik Tse Agency"
|
|
MODEL_NAME = "ChatPBC V4 Advanced" # Will be adjusted for V3.3
|
|
SYSTEM_PROMPT = f"You are ChatPBC, an expert AI business strategist developed by {DEVELOPER}. Provide strategic, actionable business advice."
|
|
|
|
# Configuration
|
|
MODEL_ID = os.environ.get("MODEL_ID", "chatpbc1/chatpbc-v4")
|
|
HF_TOKEN = os.environ.get("HF_TOKEN")
|
|
|
|
print(f"Loading model {MODEL_ID}...")
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID,
|
|
torch_dtype=torch.float16,
|
|
device_map="auto",
|
|
token=HF_TOKEN
|
|
)
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
|
|
def chat_function(message, history):
|
|
# Handle file content if present (Gradio 4+ handles this in message dict)
|
|
text = message["text"] if isinstance(message, dict) else message
|
|
files = message["files"] if isinstance(message, dict) and "files" in message else []
|
|
|
|
file_content = ""
|
|
if files:
|
|
for f in files:
|
|
try:
|
|
with open(f, "r", errors="ignore") as file:
|
|
file_content += f"\n[File: {os.path.basename(f)}]\n{file.read()[:2000]}\n"
|
|
except: pass
|
|
|
|
full_user_msg = text + (f"\n\nContext from uploaded files:\n{file_content}" if file_content else "")
|
|
|
|
# Format Llama-2 chat prompt
|
|
prompt = f"[INST] <<SYS>>\n{SYSTEM_PROMPT}\n<</SYS>>\n\n"
|
|
for user_msg, assistant_msg in history:
|
|
# history elements can be strings or dicts in newer Gradio
|
|
u = user_msg["text"] if isinstance(user_msg, dict) else user_msg
|
|
a = assistant_msg["text"] if isinstance(assistant_msg, dict) else assistant_msg
|
|
prompt += f"{u} [/INST] {a} <s>[INST] "
|
|
prompt += f"{full_user_msg} [/INST]"
|
|
|
|
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
|
|
|
|
with torch.no_grad():
|
|
output = model.generate(
|
|
**inputs,
|
|
max_new_tokens=512,
|
|
temperature=0.7,
|
|
do_sample=True,
|
|
repetition_penalty=1.1,
|
|
eos_token_id=tokenizer.eos_token_id
|
|
)
|
|
|
|
response = tokenizer.decode(output[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
|
|
return response
|
|
|
|
demo = gr.ChatInterface(
|
|
fn=chat_function,
|
|
title=MODEL_NAME,
|
|
description=f"Expert AI Business Strategist developed by {DEVELOPER}",
|
|
multimodal=True,
|
|
theme="soft"
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
demo.launch()
|