File size: 1,165 Bytes
0c16a61
593f71b
924ec7c
333d2c3
 
 
 
3b29212
333d2c3
3c63a2d
5b4fcc2
3c63a2d
3b29212
333d2c3
 
3b29212
333d2c3
3b29212
333d2c3
3b29212
333d2c3
 
3b29212
333d2c3
3b29212
333d2c3
 
 
 
3b29212
333d2c3
3b29212
333d2c3
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import gradio as gr
import os
from huggingface_hub import login
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from langchain_community.llms import HuggingFacePipeline
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

# Load model and tokenizer (Gemma 2B or similar)
model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
login(token=os.environ["HUGGINGFACE_TOKEN"])


tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype="auto")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)

llm = HuggingFacePipeline(pipeline=pipe)

# Simple prompt template
prompt = PromptTemplate.from_template("You are Krish, a wise and witty friend.\n\nUser: {question}\nKrish:")

chain = LLMChain(prompt=prompt, llm=llm)

# Gradio interface
def chat_fn(message):
    response = chain.run({"question": message})
    return response.strip()

iface = gr.Interface(fn=chat_fn, inputs="text", outputs="text", title="🦚 Meet Krish", description="A wise, witty, and compassionate friend - KrishWay")

iface.launch()