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

model_id = "mistralai/Mistral-7B-Instruct-v0.3"

# Load tokenizer and model with correct settings
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)  # Important fix
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto"
)

# Create generation pipeline
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512,
    do_sample=True,
    top_k=50,
    top_p=0.95,
    temperature=0.7,
    repetition_penalty=1.1
)

# Define Gradio UI
def chat_fn(message, history):
    prompt = f"[INST] {message.strip()} [/INST]"
    output = pipe(prompt)[0]['generated_text']
    return output.replace(prompt, "").strip()

chatbot = gr.ChatInterface(
    fn=chat_fn,
    title="🤖 Vynix AI - Powered by Mistral",
    description="Ask anything! Built using Mistral-7B-Instruct-v0.3.",
)

# Launch the app
chatbot.launch()