Vynix.Unity / app.py
Zubiiiiiii294's picture
Update app.py
22f1de6 verified
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()