Spaces:
Running
Running
File size: 2,677 Bytes
097081a 6705397 ea914de 9b56ad1 097081a 6705397 9b56ad1 6705397 3f106f4 6705397 ea914de 6705397 ea914de 6705397 ea914de 6705397 ea914de 6705397 9b56ad1 6705397 ea914de 6705397 9b56ad1 6705397 9b56ad1 6705397 |
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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import os
import time
import torch
import gradio as gr
from huggingface_hub import login
from transformers import AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM
from sentence_transformers import SentenceTransformer
from langchain_community.vectorstores import FAISS
# Load HF token and login
hf_token = os.environ.get("HUGGINGFACE_TOKEN")
if not hf_token:
raise ValueError("Please set the HUGGINGFACE_TOKEN environment variable")
login(token=hf_token)
# Load tokenizer and quantized model
model_id = "TheBloke/mistral-7B-GPTQ"
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
print("Loading quantized model...")
start = time.time()
model = AutoGPTQForCausalLM.from_quantized(
model_id,
use_safetensors=True,
device=device,
use_triton=True,
quantize_config=None,
)
print(f"Model loaded in {time.time() - start:.2f} seconds on {device}")
# Load embedding model for FAISS vector store
embedder = SentenceTransformer("all-MiniLM-L6-v2")
# Sample documents to build vector index (can replace with your own)
texts = [
"Hello world",
"Mistral 7B is a powerful language model",
"Langchain and FAISS make vector search easy",
"This is a test document for vector search",
]
embeddings = embedder.encode(texts)
faiss_index = FAISS.from_embeddings(embeddings, texts)
# Generate text from prompt
def generate_text(prompt, max_length=128):
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model.generate(**inputs, max_length=max_length)
decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
return decoded
# Search docs with vector similarity
def search_docs(query):
query_emb = embedder.encode([query])
results = faiss_index.similarity_search_by_vector(query_emb[0], k=3)
return "\n\n".join(results)
# Gradio UI
with gr.Blocks() as demo:
gr.Markdown("# Mistral GPTQ + FAISS Vector Search Demo")
with gr.Tab("Text Generation"):
prompt_input = gr.Textbox(label="Enter prompt", lines=3)
generate_btn = gr.Button("Generate")
output_text = gr.Textbox(label="Output", lines=6)
generate_btn.click(fn=generate_text, inputs=prompt_input, outputs=output_text)
with gr.Tab("Vector Search"):
query_input = gr.Textbox(label="Enter search query", lines=2)
search_btn = gr.Button("Search")
search_output = gr.Textbox(label="Search Results", lines=6)
search_btn.click(fn=search_docs, inputs=query_input, outputs=search_output)
if __name__ == "__main__":
demo.launch()
|