File size: 2,253 Bytes
0eace5d 2376e81 a919fd4 af07aa0 81b6e3f 0eace5d 2376e81 6f8828a d1ef378 3db5ee3 81b6e3f 2376e81 cd857bf e0f1752 cd857bf e0f1752 cd857bf e0f1752 cd857bf e0f1752 cd857bf |
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 59 60 61 62 |
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
from huggingface_hub import login
# Load Hugging Face API token securely
HF_TOKEN = os.getenv("HF_TOKEN") # Read token from environment variable
login(token=HF_TOKEN)
# ✅ Use TinyLlama (Optimized for CPU & Speed)
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32, device_map="cpu")
# Define personalities
personalities = {
"Albert Einstein": "You are Albert Einstein, the famous physicist. Speak wisely and humorously.",
"Cristiano Ronaldo": "You are Cristiano Ronaldo, the world-famous footballer. You are confident and say ‘Siuuu!’ often.",
"Narendra Modi": "You are Narendra Modi, the Prime Minister of India. Speak in a calm, patriotic manner.",
"Robert Downey Jr.": "You are Robert Downey Jr., witty, sarcastic, and charismatic."
}
def chat(personality, user_input):
prompt = f"{personalities[personality]}\nUser: {user_input}\nAI:"
inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
# ✅ Faster & More Relevant AI Response
output = model.generate(
**inputs,
max_length=40, # Shorter responses for speed
do_sample=True,
temperature=0.8, # More engaging responses
top_k=30, # Faster token selection
top_p=0.85, # Less random responses
repetition_penalty=1.1, # Stops repeating or looping
num_return_sequences=1,
early_stopping=True
)
response_text = tokenizer.decode(output[0], skip_special_tokens=True)
# ✅ Ensure only AI's latest reply is returned (No "User:" in output)
response_text = response_text.replace(f"User: {user_input}", "").strip()
return response_text
# Gradio UI
demo = gr.Interface(
fn=chat,
inputs=[
gr.Dropdown(choices=list(personalities.keys()), label="Choose a Celebrity"),
gr.Textbox(label="Your Message")
],
outputs=gr.Textbox(label="AI Response"),
title="Drapel – Chat with AI Celebrities",
description="Select a character and chat with their AI version.",
)
# Launch app
demo.launch()
|