Nicolai Berk commited on
Commit
ddc705c
·
1 Parent(s): e922deb

Revert to old model, login with secret token

Browse files
Files changed (1) hide show
  1. app.py +11 -2
app.py CHANGED
@@ -6,6 +6,15 @@ import torch
6
  import numpy as np
7
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
8
 
 
 
 
 
 
 
 
 
 
9
  # Load corpus
10
  print("Loading dataset...")
11
  dataset = load_dataset("rag-datasets/rag-mini-wikipedia", "text-corpus")
@@ -25,8 +34,8 @@ index.add(corpus_embeddings_np)
25
  reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
26
 
27
  # Generator (choose one: local HF model or OpenAI)
28
- tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-4")
29
- model = AutoModelForCausalLM.from_pretrained("microsoft/phi-4", device_map="auto", torch_dtype=torch.float16)
30
  generator = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150)
31
 
32
  def rag_pipeline(query):
 
6
  import numpy as np
7
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
8
 
9
+ import os
10
+
11
+ # Login automatically if HF_TOKEN is present
12
+ hf_token = os.getenv("HF_TOKEN")
13
+ if hf_token:
14
+ from huggingface_hub import login
15
+ login(token=hf_token)
16
+
17
+
18
  # Load corpus
19
  print("Loading dataset...")
20
  dataset = load_dataset("rag-datasets/rag-mini-wikipedia", "text-corpus")
 
34
  reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
35
 
36
  # Generator (choose one: local HF model or OpenAI)
37
+ tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
38
+ model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", device_map="auto", torch_dtype=torch.float16)
39
  generator = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150)
40
 
41
  def rag_pipeline(query):