app.py

import gradio as gr
import regex as re
from tqdm import tqdm
import pickle
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    pipeline,
)
import torch

def check_check():

    model_name = "NousResearch/Llama-2-7b-chat-hf"
    
    print('model loading......')
    model = AutoModelForCausalLM.from_pretrained(model_name)

    print("model loaded")
    model.config.use_cache = False
    
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    
    file_path = 'data.pkl'
    print("opening file")
    with open(file_path, 'rb') as files:
        model.state_dict = pickle.load(files)
    print("Lets go baby")
        
    # pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=1500)
    # result = pipe(text)
    # print(result[0]['generated_text'])


title = "Fine tuning Llama2-7B"
description = "A simple Gradio interface to do some shit"


def inference(text):

    return "A", "B", "C"

iface = gr.Interface(
    inference,
    inputs=["text"],
    outputs=["text", "text", "text"],
    examples=["سفید رنگ ہیں آخر سیاہ مو کرتے لٹاتے دولت دنیا کو میکدے میں ہم طلائی ساغر مے نقرئی سبو کرتے ہمیشہ میں نے گریباں کو چاک چاک کیا",
              " دل کہ آتے ہیں جس کو دھیان بہت خود بھی آتا ہے اپنے دھیان میں کیاوہ ملے تو یہ پوچھنا ہے مجھےاب بھی ہوں میں تری امان میں کیا"],
    title = title,
    description = description
)

check_check()
iface.launch()