import gradio as gr
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

# Koristimo manji i brži model
MODEL_NAME = "facebook/opt-1.3b"

# Učitavanje modela i tokenizatora sa optimizacijama
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32, device_map="auto")

# Kreiranje pipeline-a za generisanje teksta
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, torch_dtype=torch.float32, device_map="auto")

# Funkcija za generisanje odgovora
def respond(message, history=None):
    prompt = f"<s>[INST] {message} [/INST]"
    outputs = pipe(prompt, max_new_tokens=50, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
    return outputs[0]["generated_text"]

# Pokretanje Gradio chat interfejsa
gr.ChatInterface(respond).launch()