Nicolai Berk commited on
Commit
e9dd011
·
1 Parent(s): cd4d813

Switch to open model

Browse files
Files changed (1) hide show
  1. app.py +2 -3
app.py CHANGED
@@ -6,7 +6,6 @@ import torch
6
  import numpy as np
7
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
8
 
9
-
10
  # Load corpus
11
  print("Loading dataset...")
12
  dataset = load_dataset("rag-datasets/rag-mini-wikipedia", "text-corpus")
@@ -26,8 +25,8 @@ index.add(corpus_embeddings_np)
26
  reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
27
 
28
  # Generator (choose one: local HF model or OpenAI)
29
- tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
30
- model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", device_map="auto", torch_dtype=torch.float16)
31
  generator = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150)
32
 
33
  def rag_pipeline(query):
 
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
  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):