|
import gradio as gr |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
import torch |
|
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
|
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" |
|
model = AutoModelForCausalLM.from_pretrained( |
|
model_name, |
|
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, |
|
device_map="auto" if torch.cuda.is_available() else None, |
|
trust_remote_code=True |
|
).eval() |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) |
|
tokenizer.pad_token = tokenizer.eos_token |
|
|
|
|
|
def chatbot(user_input): |
|
if not user_input.strip(): |
|
return "Please enter a message." |
|
|
|
inputs = tokenizer(user_input, return_tensors="pt", padding=True).to(device) |
|
|
|
with torch.no_grad(): |
|
output = model.generate( |
|
input_ids=inputs["input_ids"], |
|
attention_mask=inputs["attention_mask"], |
|
max_new_tokens=50, |
|
temperature=0.6, |
|
top_p=0.8, |
|
do_sample=True, |
|
early_stopping=True, |
|
pad_token_id=tokenizer.eos_token_id |
|
) |
|
|
|
response = tokenizer.decode(output[0], skip_special_tokens=True) |
|
return response |
|
|
|
|
|
iface = gr.Interface(fn=chatbot, inputs="text", outputs="text", title="Octagon 2.0 Chatbot") |
|
|
|
|
|
iface.launch(ssl=False, debug=True) |