# -*- coding: utf-8 -*-
"""app.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1va0D7B5cFROqIv6pNnXrWA4-sMWQp_F2
"""

# %pip install torch torchversion torchaudio

# !pip install langchain einop accelerate transformers bitsandbytes

# !pip install langchain-community

from langchain import HuggingFacePipeline
from langchain import PromptTemplate, LLMChain
from transformers import AutoTokenizer, AutoModelForCausalLM
import transformers
import os
import torch

# torch.cuda.is_available()

model_id = "tiiuae/falcon-7b-instruct"

tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(model_id,
                                             cache_dir='./workspace/',
                                             torch_dtype=torch.bfloat16,
                                             trust_remote_code=True,
                                             device_map="auto",
                                             offload_folder="offload")
model.eval()

pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    max_length=2048,
    do_sample=True,
    top_k=10,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
    device_map="auto",
)

pipeline("what is a pizza?")

template = PromptTemplate(input_variables=['input'], template='{input}')

llm = HuggingFacePipeline(pipeline=pipeline)

llm_chain = LLMChain(prompt=template, llm=llm)

response = llm_chain.run("what is a pizza?")

print(response)

# !pip install gradio

import gradio as gr

def generate(promt):
  return llm_chain.run(promt)

title = 'GYANA.AI'
description = 'Gyana.AI is based on HKRM model it will take 4-5min to ans'
gr.Interface(fn=generate, inputs=['text'], outputs=['text'], title=title, description=description, theme='gstaff/xkcd').launch(share=True)