Updated imports in script

streaming_inference.py

@@ -3,8 +3,8 @@ from threading import Thread
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 
-from .configuration_phi import PhiConfig
-from .modeling_phi import PhiForCausalLM
+from .phi2_configuration import Phi2Config
+from .phi2_model import Phi2ModelForCausalLM
 
 
 # This works, but is not streaming
@@ -12,8 +12,8 @@ from .modeling_phi import PhiForCausalLM
 if __name__ == "__main__":
     device = "cuda"
 
-    model_config = PhiConfig(**json.load(open("simplified_phi2/config.json")))
-    model = PhiForCausalLM(model_config).to(device)
+    model_config = Phi2Config(**json.load(open("simplified_phi2/config.json")))
+    model = Phi2ModelForCausalLM(model_config).to(device)
     phi_model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", trust_remote_code=True)
     model.load_state_dict(phi_model.state_dict())
 
@@ -45,8 +45,8 @@ if __name__ == "__main__":
 
     # make model and run model.generate(streamer=TextIteratorStreamer) on a thread
     device = "cuda"
-    model_config = PhiConfig(**json.load(open("simplified_phi2/config.json")))
-    model = PhiForCausalLM(model_config).to(device)
+    model_config = Phi2Config(**json.load(open("simplified_phi2/config.json")))
+    model = Phi2ModelForCausalLM(model_config).to(device)
     phi_model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", trust_remote_code=True)
     model.load_state_dict(phi_model.state_dict())
     thread = Thread(